Skip to content

Latest commit

 

History

History

yolo-fastestv2

Folders and files

NameName
Last commit message
Last commit date

parent directory

..
 
 
 
 
 
 
 
 

yolo-fastestv2

Input --> Output

Convert

pt --> TorchScript --> pnnx --> ncnn

import torch
import model.detector   # 导入模型
import os

# model = model.detector.Detector(cfg["classes"], cfg["anchor_num"], True, True)
model = model.detector.Detector(80, 3, True, True)
model.load_state_dict(torch.load(r"modelzoo\coco2017-0.241078ap-model.pth"))
# sets the module in eval node
model.eval()


# # 1. pt --> torchscript
# traced_script_module = torch.jit.trace(model, torch.randn(1, 3, 352, 352))
# traced_script_module.save("ts.pt")
# # 2. ts --> pnnx --> ncnn
# os.system("pnnx ts.pt inputshape=[1,3,352,352]")

torch.onnx.export(model,  # model being run
                  torch.randn(1, 3, 352, 352),                 # model input (or a tuple for multiple inputs)
                  "out.onnx",               # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=11,          # the ONNX version to export the model to
                  do_constant_folding=True)  # whether to execute constant folding for optimization

# 2. onnx --> onnxsim
os.system("python -m onnxsim out.onnx sim.onnx")

# 3. onnx --> ncnn
os.system("onnx2ncnn sim.onnx ncnn.param ncnn.bin")

# 4. ncnn --> optmize ---> ncnn
os.system("ncnnoptimize ncnn.param ncnn.bin opt.param opt.bin 1")  # 数字0 代表fp32 ;1代表fp16

转换方法,将上述文件放到项目目录下,然后运行命令就可以成功。 torchscript转换成功,但是运行失败,onnx方法转换成功

Example project

Reference