yolact转onnx

利用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)