PyTorch YOLOV3 模型转换问题

工程代码：https://github.com/eriklindernoren/PyTorch-YOLOv3

第一步：下载Darknet 版的yolov3.weights

第二部：使用以上工程代码加载模型：model.load_darknet_weights() 并保存为pytorch模型格式yolov3.pth: 

torch.save(model.state_dict(), pytorch_model)

第三步：用 torch.jit.trace   convert  模型为  .pt  格式。 这里注意：参数check_trace=False，不然会报错

traced_script_module = torch.jit.trace(model, example,check_trace=False)

参考：https://github.com/pytorch/pytorch/issues/23993

注意：torch.jit.trace（）不支持控制流语句，需将models.py 中 if targets is None 以下语句注释掉，增加

return output,0 因为作为推断，所以注释掉的语句不影响。

import numpy as np
import torch
# import torchsnooper
from models import Darknet
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# with torch.jit.optimized_execution(True):

def save_pytorch_model(darknet_model,pytorch_model,config_file):

    model = Darknet(config_file).to(device)
    model.load_darknet_weights(darknet_model)
    torch.save(model.state_dict(), pytorch_model)

    print('pytorch model saved!')

def convert_pytorch_model_to_libtorch(config_file,pytorch_model,libtorch_model):

    model = Darknet(config_file).to(device)

    model.load_state_dict(torch.load(pytorch_model))

    model.eval()
    example = torch.rand(1,3,416,416).cuda()
    with torch.jit.optimized_execution(True):
        print('pp')
        # example 报错 Expected type 'tuple', got 'Tensor' instead ，可能是input没有放cuda上
        traced_script_module = torch.jit.trace(model, example,check_trace=False)# save the converted model
        print('oo')
        traced_script_module.save(libtorch_model)

        output = traced_script_module(torch.rand(1,3,416,416).cuda())
        print(output)

    print('trace model saved!')