1. 程式人生 > 其它 >pytorch模型轉NCNN模型在手機部署

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