Yolov3(Mxnet)测试加速：GPU图像预处理

深度学习的数据标准化操作在测试的也要遵守，但是Mxnet中Gluoncv使用CPU的串行数据标准化，对于某些实时性要求较高的任务，在CPU使用率较高时，数据标准化的耗时严重拖累了网络的预测速度。我们以Mxnet中的Yolov3为例，介绍使用GPU进行标准化的方法，减少CPU负担。首先需要下载Gluoncv代码，并将其中的模型增加一步卷积操作用作预处理，这样模型便可以直接在GPU上处理原图。除了可以直接在GPU上使用broadcast_sub和broadcast_div进行标准化处理之外，还可以使用一步卷积实现。

Gluoncv代码链接：https://github.com/dmlc/gluon-cv

图像预处理通常使用imagenet上的标准预处理方法，归一化到0-1之间，减去均值再除以标准差。通常使用的mean和std为：

mean=(0.485, 0.456, 0.406)
std=(0.229, 0.224, 0.225)

计算方法大致可以等效为以下公式：

将公式中 作为weight， 作为bias，预处理过程就可以看作是1*1卷积。因此图像的预处理或许可以使用卷积来实现，并拼接到训练好的网络上，实现端到端的测试。由于RGB三个通道分开处理，所以可以使用Group卷积，3个1x1的卷积核，3个bias，使用上述参数值初始化weight和bias。实际上不使用卷积，使用普通的broadcast加减乘除也可以在GPU实现这个操作，然后传给网络，这个比较简单，这里就不再详细介绍。

1. 更改网络模型

首先更改gluoncv/model_zoo/yolo/yolo3.py中的YOLOV3类，在__init__函数中的增加一个预处理层preprocess：

with self.name_scope():
    self.preprocess = nn.HybridSequential()
    self.preprocess.add(nn.Conv2D(channels=3, kernel_size=1, strides=1, padding=0, groups=3, use_bias=True))

然后便可以在hybrid_forward函数最开始部分使用如下代码完成数据预处理，网络便可以直接接受Mxnet中的读图mx.image.imread的结果,前三步好像省略不了，这也是预处理中必须实现的步骤：

x = x.astype('float32') #to float32
x = F.transpose(x, axes=(2, 0, 1)) #h*w*c to c*h*w
x = x.expand_dims(0) #c*h*w to 1*c*h*w
x = self.preprocess(x)

因为网络中新增了一个带参数的层，因此加载默认的VOC或者COCO预训练模型会报错，因此将yolo3.py中的get_yolov3中的load_parameters增加参数allow_missing=True：

net.load_parameters(get_model_file(full_name, tag=pretrained, root=root), ctx=ctx, allow_missing=True)

至此，我们可以将修改gluoncv文件夹重命名为gluoncv_pre，然后使用import导入模型。

2. 改写测试程序

原始测试程序demo_yolo.py链接：https://gluon-cv.mxnet.io/build/examples_detection/demo_yolo.html，测试程序需要导入修改的模型：

from gluoncv_pre import model_zoo

然后通过get_model得到预训练模型：

net = model_zoo.get_model('yolo3_darknet53_voc', pretrained=True)

至此加载的模型包含了我们增加的preprocess层之外的所有参数，而我们增加的层需要自己初始化。首先定义初始化方法，分别将weight和bias按照三个通道不同，进行不同初始化：

class Preprocess_weightInit(mx.init.Initializer):
    def __init__(self):
        super(Preprocess_weightInit, self).__init__()
    def _init_weight(self, _, arr):
        arr[0] = 1.0 / float(255) / 0.229
        arr[1] = 1.0 / float(255) / 0.224
        arr[2] = 1.0 / float(255) / 0.225

class Preprocess_biasInit(mx.init.Initializer):
    def __init__(self):
        super(Preprocess_biasInit, self).__init__()
    def _init_weight(self, _, arr):
        arr[0] = -0.485 / 0.229
        arr[1] = -0.456 / 0.224
        arr[2] = -0.406 / 0.225

然后使用上述初始化方法初始化我们preprocess层参数：

net.preprocess.collect_params()['yolov30_conv0_weight'].initialize(Preprocess_weightInit())
net.preprocess.collect_params()['yolov30_conv0_bias'].initialize(Preprocess_biasInit())

至此，所有参数初始化完成，直接使用imread读取图像，经过resize到测试尺寸后，便可以直接将图像传入网络：

img_ori = mx.image.imread(im_fullpath, to_rgb=1)
img_resize = timage.resize_short_within(img_ori, 1024, max_size=1024, mult_base=1)
img_resize = img_resize.as_in_context(ctx)
output = net(img_resize)

实际上，GPU上也可是实现resize，如果测试图像基本尺寸固定，也可以在网络中使用如下方式resize，不过貌似CPU resize和GPU实现方式有点区别，对精度有一些影响：

x = x.astype('float32') #int8 to float32
x = F.transpose(x, axes=(2, 0, 1)) #h*w*c to c*h*w
x = x.expand_dims(0) #c*h*w to 1*c*h*w
x = F.contrib.BilinearResize2D(data=x, height=512, width=512)
x = self.preprocess(x)

这样就可以直接imread读图传给网络，减少CPU的使用：

img_ori = mx.image.imread(im_fullpath, to_rgb=1)
img_ori = img_ori.as_in_context(ctx)
output = net(img_ori)

至此，整个网络可以按照另一篇博客中的导出方法导出成json和params文件，下次可以直接load使用。

注：不知道是否遗漏了一些预处理的细节。

------------------------------------------------------------------------------------------------------------------------------------------------------------------------

感谢weixin_39534008指出，目前Gluoncv中utils中的export_block在export外又封装了一层，支持将preprocess打包到网络前面，然后一起export；看了下代码，export_block定义了一个预处理类_DefaultPreprocess，采用broadcast_sub和broadcast_div进行标准化，使用一个HybridSequential add了_DefaultPreprocess和原始net得到wrapper_block，然后再wrapper_block.export(path, epoch)。