欢迎您访问 最编程 本站为您分享编程语言代码,编程技术文章!
您现在的位置是: 首页

yolact turn onnx

最编程 2024-03-15 09:54:33
...

利用yolact训练了一波自己的数据,目前效果还可以,准备利用C++调用,这里记录一下模型转为onnx的过程和一些问题。

一、模型转换

1、设置配置文件

依据自己训练的配置,设置配置文件

set_cfg("yolact_resnet50_custom_config")

2、精简yolact.py

依据配置参数,删除一些判断分支语句,yolact只保留初始化、加载权重和推理

3、加载权重

weight_path = "../weights/20200428_754_80000.pth"
net = Yolact()
net.load_weights(weight_path)
net.eval()
net.to(device)

4、转换

inputs = torch.randn(1, 3, 550, 550).to(device)
onnx_model_path = "./yolact.onnx"
print("convert net to ", onnx_model_path, " ... ")
torch.onnx.export(
    net,
    (inputs,),
    onnx_model_path,
    verbose=True,
    input_names=["img"],
    output_names=["loc", "conf", "mask", "proto"],
    opset_version=12
)
print("converted successed!")

二、常见问题

1、FPN不是子模块

  • 错误信息
    RuntimeError: Tried to trace <torch.yolact.FPN object at 0x9c26df60> but it is not part of the active trace. Modules that are called during a trace must be registered as submodules of the thing being traced.

  • 解决方法
    找到yolact.py文件中FPN类的位置

class FPN(ScriptModuleWrapper):

修改为

class FPN(nn.Module):

2、device报错

  • 错误信息
    RuntimeError: legacy constructor for device type: cpu was passed device type: cuda, but device type must be: cpu.

  • 解决方法
    在yolact开头添加

device = "cuda" if torch.cuda.is_available() else "cpu"

找到yolact.py文件中

self.priors = torch.Tensor(prior_data, device=device).view(-1, 4).detach()

将其更改为

self.priors = torch.Tensor(prior_data).view(-1, 4).detach().to(device)

原文地址:https://www.cnblogs.com/xiaxuexiaoab/p/16249384.html