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