pytorch模型转NCNN模型在手机部署

记录一下自己的转换过程，中间也踩了很多坑。。

第一步， pytorch转onnx

这一步比较方便，pyroch自身就支持，注意input和output_names一定要填写正确。

net = torch.load("./model17_4.pkl")

net.eval()

test_arr = np.random.rand(1, 3, 400, 400)
dummy_input = torch.Tensor(test_arr).cuda()

# print(dummy_input.shape)
# d = dummy_input.permute(0, 3, 1, 2)
# print(d.shape)

input_names = ["input"]
output_names = ["output"]
torch.onnx.export(net,
dummy_input,
"mobilenet_v7_2.onnx",
export_params=True,
do_constant_folding=True,
input_names=input_names,
output_names=output_names)

得到onnx模型后，还可以写代码验证一下效果：

TF_PATH = "tf_model6" # where the representation of tensorflow model will be stored
ONNX_PATH = "mobilenet_v7_2.onnx" # path to my existing ONNX model
onnx_model = onnx.load(ONNX_PATH) # load onnx model
tf_rep = prepare(onnx_model) # creating TensorflowRep object
tf_rep.export_graph(TF_PATH)

 

#test tensorflow mode

img = cv2.imread('./data/000701.jpg')[:,:,::-1]
#img = img[850:1250,850:1250]

img = (np.transpose(img, (2, 0, 1)) - 128.) / 128.

#img = torch.from_numpy(img).unsqueeze(0).float()
#input = img.numpy()
input = tf.expand_dims(img, 0).numpy()

output = tf_rep.run(input) # run the loaded model
res = output.output[0]
res = np.clip(res, 0, 1)
im = np.array(res*255, dtype=np.uint8)
im1 = np.transpose(im, (1, 2, 0))[:,:,::-1]
cv2.imwrite('./test/out.bmp', im1)

第二步，将onnx模型转换为ncnn模型

这个参考了这篇文章，还不错https://zhuanlan.zhihu.com/p/137458205

先到github下载最新的ncnn库，然后编译。注意需要安装opencv和protobuf，否则编译ncnn会不完整或者有异常。

编译成功后，会在build/tools/目录下生成onnx和quantize等文件夹，

cd进入onnx文件夹，把上一步生成的.onnx文件拷贝过来，

先简化模型：

python3 -m onnxsim mobilenet_v7_2.onnx  mobilenet_v7_2.onnx-sim.onnx

再运行

./onnx2ncnn mobilenet_v7_2.onnx-sim.onnx  model.param model.bin

 

生成的.bin与.param文件就是我们在安卓上需要使用的NCNN模型文件

第三步，在手机上加载运行ncnn模型

 

ncnn::Net net;
//net.opt.num_threads=1;
net.load_param("./model.param");
net.load_model("./model.bin");
// 把opencv的mat转换成ncnn的mat
ncnn::Mat input = ncnn::Mat::from_pixels(img2.data, ncnn::Mat::PIXEL_BGR, img2.cols, img2.rows);
const float mean_vals[3] = { 128.f,128.f,128.f };
const float norm_vals[3] = { 1 / 128.f,1 / 128.f,1 / 128.f };
input.substract_mean_normalize(mean_vals, norm_vals);
// ncnn前向计算
ncnn::Extractor extractor = net.create_extractor();
extractor.input("input", input);
ncnn::Mat m, output1;//取决于模型的输出有几个

clock_t start_time = clock();

extractor.extract("output", output1);

 

最后，再记录一下我的相关软件版本，注意onnx-tf的版本不能太低：

tensorflow-gpu 2.8.0

onnx 1.6.0

onnx-tf 1.6.0

libprotoc 3.5.0