Github MegEngine models 有现成的 imagenet 预训模型。这里把模型 dump 成 .mge。
新增 dump.py,按 [1, 3, 224, 224] 尺寸 trace 模型,打开推理优化选项,保存为 model.mge
。
$ git clone https://github.com/MegEngine/models
$ cd models
$ export PYHTONPATH=${PWD}:${PYTHONPATH}
$ cd official/vision/classification/resnet
$ python3 dump.py
$ ls -lah model.mge
...
dump.py
已经 PR 到 MegEngine/models 分类模型目录
$ cat dump.py
...
data = mge.Tensor(np.random.random((1, 3, 224, 224))) # 准备一个样例输入
@jit.trace(capture_as_const=True)
def pred_func(data):
outputs = model(data) # trace 每个 opr 的 shape
return outputs
pred_func(data)
pred_func.dump( # 保存模型
graph_name,
arg_names=["data"],
optimize_for_inference=True, # 打开推理优化选项
enable_fuse_conv_bias_nonlinearity=True, # 打开 fuse conv+bias+ReLU pass 推理更快
)
...
开发模型的推理封装,对外提供功能内聚的接口。调用方传入一张或多张图片、直接获取结果,尽量避免关心内部实现(如用何种 backbone、预处理是什么、后处理是什么)。
$ cat flow-python/examples/application/simple_classification/lite.py
...
def inference(self, mat):
img = self.preprocess(mat, input_size=(224,224), scale_im = False, mean=[103.530, 116.280, 123.675], std=[57.375, 57.120, 58.395])
# 设置输入
inp_data =self.net.get_io_tensor("data")
inp_data.set_data_by_share(img)
# 推理
self.net.forward()
self.net.wait()
# 取输出
output_keys = self.net.get_all_output_name()
output = self.net.get_io_tensor(output_keys[0]).to_numpy()
return np.argmax(output[0])
...
$ python3 lite.py --model model.mge --path test.jpg # 测试
2021-09-14 11:45:02.406 | INFO | __main__:<module>:81 - 285
285
是分类模型最后一层的 argmax,对应含义需要查 imagenet 数据集分类表 ,这里是 “Egyptian cat”(下标从 0 开始)。
MegEngine 除了不需要转模型,还能消除预处理。我们修改 dump_resnet.py 把预处理从 SDK/业务代码提到模型内。这样的好处是:划清工程和算法的边界,预处理本来就应该由 scientist 维护,每次只需要 release mge 文件,减少交接内容
...
@jit.trace(capture_as_const=True)
def pred_func(data):
out = data.astype(np.float32)
output_h, output_w = 224, 224
# resize
print(shape)
M = mge.tensor(np.array([[1,0,0], [0,1,0], [0,0,1]], dtype=np.float32))
M_shape = F.concat([data.shape[0],M.shape])
M = F.broadcast_to(M, M_shape)
out = F.vision.warp_perspective(out, M, (output_h, output_w), format='NHWC')
# mean
_mean = mge.Tensor(np.array([103.530, 116.280, 123.675], dtype=np.float32))
out = F.sub(out, _mean)
# div
_div = mge.Tensor(np.array([57.375, 57.120, 58.395], dtype=np.float32))
out = F.div(out, _div)
# dimshuffile
out = F.transpose(out, (0,3,1,2))
outputs = model(out)
return outputs
...
具体实现是在 trace inference 里增加预处理动作,fuse opr 优化加速的事情交给 MegEngine 即可。更多 cv 操作参照 MegEngine API 文档。
因为推理输入变成了 BGR,所以 dump 模型的时候参数也应该跟着变
$ python3 dump.py -a resnet18 -s 1 224 224 3