从人工智能鉴黄模型，尝试TensorRT优化

随着互联网的快速发展，愈来愈多的图片和视频出如今网络，特别是UCG产品，激发人们上传图片和视频的热情，好比微信天天上传的图片就高达10亿多张。每一个人均可以上传，这就带来监管问题，若是没有内容审核，色情图片和视频就会泛滥。前不久，一贯以开放著称的tumblr，就迫于压力，开始限制人们分享色情图片。更别提国内，内容审核是UCG绕不过去的坎。还记得前几年出现的职业鉴黄师这一职业么？传说百万年薪，天天看黄片看得想吐，但最近又不多有人说起这一职业，这个应监管而生的职业，因人工智能的出现又快速消亡。（固然也不是彻底消亡，毕竟判断是否色情图片是一个主观的事情，有些艺术和色情之间的边界比较模糊，须要人工加以判断）

以前写过一篇文章利用人工智能检测色情图片，也曾经尝试过在浏览器中加入色情图片过滤功能，但实验下来，推理速度太慢（当时使用的Google Nexus 4作的测试，检测一张图片须要几秒钟），无法作实时过滤。最近在研究nvidia的Jetson Nano以及推理加速框架TensorRT，所以想尝试一下，看可否应用一些加速方法，加速推理。

虽然个人最终目标是应用到Jetson Nano，可是TensorRT其实适用于几乎全部的Nvidia显卡，为了方便起见，我仍是先在PC端进行尝试。没有Nvidia显卡？也没有关系，能够看看我前面发布的两篇文章：

1. 谷歌GPU云计算平台，免费又好用
2. Google Colab上安装TensorRT

open_nsfw

本文采用的深度学习模型是雅虎开源的深度学习色情图片检测模型open_nsfw，这里的NSFW表明Not Suitable for Work，该项目基于caffe框架。因为我主要研究的是Tensorflow，因此在网上找到该模型的Tensorflow实现版本，fork了一份，并添加了TensorRT框架的处理脚本，你可使用以下命令得到相关代码：

git clone https://github.com/mogoweb/tensorflow-open_nsfw.git
复制代码

在 model.py 中，咱们能够看到open_nsfw的模型定义，data/open_nsfw-weights.npy 是采用工具从yahoo open_nsfw的cafee权重转换获得的Tensorflow权重，这样咱们无需训练模型，直接用于推理过程。classify_nsfw.py 脚本可用于单张图片的推理：

python classify_nsfw.py -m data/open_nsfw-weights.npy test.jpg
复制代码

注意：脚本提供了两种解码图片文件的方式，一种是采用PIL.image、skimage进行图片处理，也就是所谓的yahoo_image_loader，一种是采用tensorflow中的图片处理函数进行处理。由于原始的open_nsfw模型是采用PIL.image、skimage进行预处理而训练的，而不一样的库解码出来的结果存在细微的差别，会影响最终结果，通常优选选择yahoo_image_loader。固然，若是你打算本身训练模型，那选择哪一种图片处理库均可以。

tools 目录下有一些脚本，能够将模型导出为frozen graph、saved model以及tflite等格式，这样咱们能够方便的在服务器端部署，还能够应用到手机端。

opt是我编写的采用TensorRT框架加速的代码，在下面我将详细说明。

导出为TensorRT模型

目前TensorRT做为Tensorflow的一部分获得Google官方支持，其包位于tensorflow.contrib.tensorrt，在代码中加入：

import tensorflow.contrib.tensorrt as trt
复制代码

就可使用TensorRT，由于有Google的支持，导出到TensorRT也就至关简单：

trt_graph = trt.create_inference_graph(
                input_graph_def=frozen_graph_def,
                outputs=[output_node_name],
                max_batch_size=1,
                max_workspace_size_bytes=1 << 25,
                precision_mode='FP16',
                minimum_segment_size=50
        )

        graph_io.write_graph(trt_graph, export_base_path, 'trt_' + graph_name, as_text=False)
复制代码

其中：

• input_graph_def 为须要导出的Tensorflow模型图定义
• outputs 为输出节点名称
• max_batch_size 为最大的batch size限制，由于GPU存在显存限制，须要根据GPU memory大小决定，通常状况能够给8或者16
• precision_mode 为模型精度，有FP3二、FP16和INT8可选，精度越高，推理速度越慢，也要依GPU而定。

graph_io.write_graph 将图写入到文件，在后续的代码中能够加载之。

完整的代码请参考 opt/export_trt.py 文件。

测试数据

由于一些政策法规的限制，并无公开数据集可提供下载，不过在github上有一些开源项目，提供脚本，从网络上进行下载。我使用的是 github.com/alexkimxyz/… 这个开源项目中的脚本。这个项目提供drawings、hentai、neutral、porn、sexy四种类别图片，能够划分为训练集和测试集，并检查图片是否有效（由于从网络爬取，有些连接不必定能访问到）。

注意这个图片下载量很是大，须要注意别把硬盘撑满。虽然这个数据量够大（几万张），能够自行进行模型训练，但和yahoo训练open_nsfw模型的图片量相比，仍是小巫见大巫，听说yahoo训练这个模型用了几百万张的图片。

推理速度对比

在opt目录下，我针对两种模型的加载和推理添加了两个脚本，分别是 benchmark_classify_nsfw.py 和 benchmark_classify_trt.py，细心的同窗可能会发现，这两个脚本几乎如出一辙，是的，除了 benchmark_classify_trt.py 多了一行代码：

import tensorflow.contrib.tensorrt as trt
复制代码

加入这行import语句，告诉tensorflow使用TensorRT框架，不然的话，会出现以下错误：

tensorflow.python.framework.errors_impl.NotFoundError: Op type not registered 'TRTEngineOp' in binary running on alex-550-279cn. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) `tf.contrib.resampler` should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed.
复制代码

取2000张测试图片进行测试，在个人GTX 960上，推理速度以下：

未优化模型： 53 s
使用TensorRT优化模型： 54 s
复制代码

若是你下载更大的数据集，能够多测试一些图片，看看优化效果。

在Google Colab上，我放了一份Jupter Notebook，有兴趣的同窗能够借助Google Colab尝试一下，文件地址：colab.research.google.com/drive/1vH-G… ，固然你也能够访问我github上完整的脚本及Notebook：

github.com/mogoweb/ten…

点击阅读原文，能够跳转到该项目。

题外话：

微信公众号流量主的门槛已经大大下降，我在公众号文章底部开通了广告，但愿没有影响你们的阅读体验。我一直很好奇，这种广告会有人点击么，过一段也许我会获得答案。

你还能够读：

1. 利用人工智能检测色情图片
2. 谷歌GPU云计算平台，免费又好用
3. Google Colab上安装TensorRT