转换 pytorch 格式模型为 caffe格式模型 pth2caffemodel

基于 GitHub xxradon/PytorchToCaffe 源码，修改 example\resnet_pytorch_2_caffe.py 如下

import os
import sys
sys.path.insert(0, '.')

import torch
from torch.autograd import Variable
from torchvision.models import resnet
import pytorch_to_caffe


"""
    resnet models in pytorch format can be downloaded from
        ‘resnet18’: ‘https://download.pytorch.org/models/resnet18-5c106cde.pth’,
        ‘resnet34’: ‘https://download.pytorch.org/models/resnet34-333f7ec4.pth’,
        ‘resnet50’: ‘https://download.pytorch.org/models/resnet50-19c8e357.pth’,
        ‘resnet101’: ‘https://download.pytorch.org/models/resnet101-5d3b4d8f.pth’,
        ‘resnet152’: ‘https://download.pytorch.org/models/resnet152-b121ed2d.pth’,

"""

def show_usage(cmd):
    print( "Usage:" )
    print(   "    ", cmd, " <pytorch-model-name>  <pytorch-model-filename.pth>" )
    
def main(cmd, argv):
    if( len(argv) < 2 ):
        print( "Error! Parameter is not enough." )
        show_usage( cmd )
        exit( 1 )

    model_name = argv[0]
    input_file = argv[1]

    pure_path = os.path.splitext( input_file )
    file_name = pure_path[0]
    
    print( " model  : ",  model_name )
    print( " input  : ",  input_file )
    print( " output : ",  '{}.prototxt'.format(file_name) )
    print( "          ",  '{}.caffemodel'.format(file_name) )
    
    
    input=torch.ones([1,3,224,224])
    match model_name:
        case "resnet18":
            resnet_x = resnet.resnet18()
        case "resnet34":
            resnet_x = resnet.resnet34()
        case "resnet50":
            resnet_x = resnet.resnet50()
        case "resnet101":
            resnet_x = resnet.resnet101()
        case "resnet152":
            resnet_x = resnet.resnet152()
        case _:
            print( "Error! Unknown model name : ",  model_name )
            show_usage( cmd )
            exit( 2 )

    if( False == os.path.isfile(input_file) ):
        print( "Error! Cannot find input file : ", input_file )
        show_usage( cmd )
        exit( 3 )

    checkpoint = torch.load(input_file)
    
    resnet_x.load_state_dict(checkpoint)
    resnet_x.eval()
    pytorch_to_caffe.trans_net(resnet_x,input,model_name)
    pytorch_to_caffe.save_prototxt('{}.prototxt'.format(file_name))
    pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format(file_name))
    
    
if __name__ == "__main__":
   main(sys.argv[0], sys.argv[1:])

脚本依赖pytorch，安装之。

pip install torch

运行中遇到 protobuf 版本过高问题，降级处理

pip install -U protobuf==3.20 

下载 resnet model文件后，执行脚本

python example\resnet_pytorch_2_caffe.py  resnet152  resnet152-b121ed2d.pth