torch模型转tensorflow模型汇总
1,针对已训练的模型,有两种思路:
(1) 通过onnx-tf模块转换:
a,具体流程:
pth=>onnx=>pb
b,环境配置:
conda create -n pth_pb python=3.7
pip install tensorflow2.1.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install tensorflow-addons0.9.1 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install onnx-tf1.5.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install onnx1.6.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
conda install pytorch torchvision #(这里我用的版本为pytorch1.4, torchvision0.5.0)
c,代码:
pth转onnx
import torchvision
import torch.onnx
import torch.nn as nn
def resnet50():
model = torchvision.models.resnet50(pretrained=False)
model.fc = nn.Linear(2048, 2)
return model
model = resnet50()
model = torch.nn.DataParallel(model)
pthfile = my.pth'
loaded_model = torch.load(pthfile, map_location='cpu')
model.load_state_dict(loaded_model['state_dict'])
input = torch.randn(1, 3, 200, 200)
input_names = ["head_input"]
output_names = ["output"]
onnx_filename = “my.onnx"
torch.onnx.export(model.module, input, onnx_filename, verbose=True, input_names=input_names, output_names=output_names)
参数 input_names表示模型的输入参数(随便起名字),output_names表示输出名字
onnx转pb