Mxnet(symbol和Gluon)并行测试

我们都知道，在训练过程中，一个batch通常包含多个训练图像，而在测试阶段，通常的做法是每次只测试一幅图像，如果我们想通过图像旋转、翻转等方法对于一幅图像进行多种输入测试时，最常想到的就是变换完之后，分别输入网络中。除此之外，还可以把多张同尺寸的图像拼成一个batch，直接一次性得到多张图像的预测结果。并行预测在某些并发较高的场景中能够显著提高系统的吞吐量。下面分别以分类模型为例，介绍symbol和Gluon接口实现方式，symbol接口实现方式很久没测试过，最近一直使用Gluon，不知symbol是否有变动。

1. symbol接口实现

通常我们加载训练好的模型和初始化的时候，需要指定data_shape，而一般测试的data_shape都是1*3*m*n，表示测试输入是一张图像。然后我们将一幅图像处理成网络可以输入的batch形式，就可以通过forward()，进行预测。其实，多个输入也只需要修改这几个地方，首先，我们修改模型加载和初始化的地方：

sym, arg_params, aux_params = mx.model.load_checkpoint('vgg', 10)
mod = mx.mod.Module(symbol=sym, context=mx.gpu(3), label_names=None)
mod.bind(for_training=False, data_shapes=[('data', (4,3,224,224))],label_shapes=mod._label_shapes)
mod.set_params(arg_params, aux_params, allow_missing=True)

我们只需要把data_shape，按照我们的输入修改，比如我们想一次输入4张图像，那么这里就可以写成（4,3,224,224），接下来我们修改生成batch的类：

class OneDataBatch():
    def __init__(self,img): 
        self.data = [mx.nd.array(img)]
        self.label = None
        self.provide_label = None
        self.provide_data = [("data",(4,3,224,224))]

接下来，我们将一张256的图像固定crop成4张224，然后将四张图拼成一个batch，并行输入网络进行测试：

for i in {8, 24}:
   for j in {8, 24}:
       im_crop = img[:, j : j + 224, i : i + 224]
       normed_img_list.append(normed_img)
im_batch = OneDataBatch(normed_img_list)
mod.forward(im_batch)
prob = mod.get_outputs()[0].asnumpy())

除了固定位置crop，我们还可以mxnet中的随机crop函数：

 for i in range(0, 4):
    cropped_im, rect  = mx.image.random_crop(mx.nd.array(im_resize), (224, 224))
    im_crop = im[:, rect[0] : rect[0] + rect[2], rect[1] : rect[1] + rect[3]]
    normed_img_list.append(normed_img)

2. Gluon接口实现

Gluon的实现更加简单，不需要复杂的配置，只要网络支持，直接将数据拼成B*3*h*w的batch就可以直接传入网络：

import mxnet as mx
from gluoncv.data.transforms.image import ten_crop, resize_short_within

transform_fn = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
    ])

img = mx.image.imread('1.jpg')
img = resize_short_within(img, short=224)
img_to_test = ten_crop(img_tree, size=(224, 224)) #crop and flip , 10*3*224*224
img_to_test = transform_fn(img_to_test)
prob = net(img_to_test)

如果是想测试多少图像，只需要使用concat函数，在B这一维度（第0维）将图像拼起来传入网络：

#img1:1*3*224*224;img2:1*3*224*224
from mxnet import nd
img_to_test = nd.concat(img1, img2, dim=0) #2*3*224*224
output = net(img_to_test)

Gluoncv中的分类模型基本都支持并行测试，而且检测模型Yolov3也支持，实际测试发现，并行测试在GPU使用率不高时，确实能提升单张图像的平均预测时间。当然并行跑两张的耗时肯定比串行跑一张的总耗时多，但是比串行跑两张的总耗时少，所以可以根据硬件水平和具体任务中等待处理的个数，合理安排并行数量。请求少时，通过翻转crop并行测试提升精度，请求多时，通过并行多张提升速度。并行测试外加在之前文章中提到的使用卷积进行预处理，能够更快，否则CPU的预处理速度会限制并行速度。

下面简单介绍Gluoncv 中的Yolov3并行测试的方法，假如我们有多少图像，而且尺寸一样（不一样的情况没测，感觉应该不行吧），不需要翻转，只需要拼接起来进行测试，class_IDs等输出的第0维便表明了是哪张图像索引。

from gluoncv import data, model_zoo
net = model_zoo.get_model('yolo3_darknet53_voc', pretrained=True)
net.collect_params().reset_ctx(mx.gpu(0))
a, img = data.transforms.presets.yolo.load_test(im1_path, short=512) #1*3*512*512
b, img = data.transforms.presets.yolo.load_test(im2_path, short=512) #1*3*512*512
img_to_test = nd.concat(a, b, dim=0)
img_to_test = img_to_test .as_in_context(mx.gpu(0))
class_IDs, scores, bounding_boxs = net(img_to_test)

假如我们有一张图像，想通过翻转进行测试提高准确率，可以通过如下方式：

from gluoncv import data, model_zoo
from mxnet import nd
net = model_zoo.get_model('yolo3_darknet53_voc', pretrained=True)
net.collect_params().reset_ctx(mx.gpu(0))

a, img = data.transforms.presets.yolo.load_test(im1_path, short=512) #1*3*512*512
#分别进行水平、垂直、对角线
x_flip_2 = nd.flip(data=x, axis=2)
x_flip_3 = nd.flip(data=x, axis=3)
x_flip_4 = nd.flip(data=x_flip_2, axis=3)

img_to_test = nd.concat(a, x_flip_2, x_flip_3, x_flip_4, dim=0)
img_to_test = img_to_test.as_in_context(mx.gpu(0))
class_IDs, scores, bounding_boxs = net(img_to_test)

最后三个输出的第0维分别存着4组数据，别忘了把box在flip回来（gluoncv/data/transforms/box.py中有对box进行flip的函数，可以import），可能还要再经过一个nms删掉重复的box。

from gluoncv.data.transforms import bbox as tbox
box_after = tbbox.flip(box_ori, (w, h), flip_x=True)