torch模型转tensorflow模型汇总

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

import onnx
from onnx_tf.backend import prepare

def onnx2pb(onnx_input_path, pb_output_path):
    onnx_model = onnx.load(onnx_input_path)  # load onnx model
    tf_exp = prepare(onnx_model)  # prepare tf representation
    tf_exp.export_graph(pb_output_path)  # export the model

if __name__ == "__main__":
    onnx_input_path = 'test.onnx'
    pb_output_path = 'test.pb'
    onnx2pb(onnx_input_path, pb_output_path)

测试pb文件

import tensorflow as tf
from torchvision import transforms
import numpy as np
from PIL import Image

transform = transforms.Compose([
    transforms.Resize((200, 200)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
with tf.Graph().as_default():
    # output_graph_def = tf.GraphDef() #tensorflow1.4版本
    output_graph_def = tf.compat.v1.GraphDef()
    output_graph_path = ‘my.pb'
    with open(output_graph_path, 'rb') as f:
        output_graph_def.ParseFromString(f.read())
        _ = tf.import_graph_def(output_graph_def, name="")
    # with tf.Session() as sess:  #tensorflow1.4版本
    with tf.compat.v1.Session() as sess:
        image = “demo.jpg"
        image_np = Image.open(image)
        img_input = transform(image_np).unsqueeze(0)
        image_np_expanded = img_input.numpy()
        # sess.run(tf.global_variables_initializer())  #tensorflow1.4版本
        sess.run(tf.compat.v1.global_variables_initializer())
        input = sess.graph.get_tensor_by_name("head_input:0")
        output = sess.graph.get_tensor_by_name("output:0")
        predictions = sess.run(output, feed_dict={input: image_np_expanded})
        index = np.argmax(predictions)
        print("predictions:", predictions)
        print("index:", index)

d，若以上流程，从onnx转pb时出现无法解决的警告，可以参考以下命令行的操作流程：
git clone https://github.com/onnx/onnx-tensorflow.git
cd onnx-tensorflow
pip install -e .
onnx-tf convert -i /Users/haobing1/myMac/model/1/resnet50_epoch_100.onnx -o /Users/haobing1/myMac/model/1/res.pb
无警告onnx转换pb
(2) 通过pytorch2keras模块转换
a，具体流程：
直接通过pytorch2keras模块将pth转成keras模型
b，环境配置：
pip3 install pytorch2keras
c，代码：

from pytorch2keras.converter import pytorch_to_keras
def converted_fully_convolutional_resnet18(
    input_tensor, pretrained_resnet=True):
    # define input tensor
    input_var = Variable(torch.FloatTensor(input_tensor))
    # get PyTorch ResNet18 model
    model_to_transfer = FullyConvolutionalResnet18(pretrained=pretrained_resnet)
    model_to_transfer.eval()
    # convert PyTorch model to Keras
    model = pytorch_to_keras(
        model_to_transfer,
        input_var,
        [input_var.shape[-3:]],
        change_ordering=True,
        verbose=False,
        name_policy="keep",
    )
    return model

其中函数pytorch_to_keras，在实现pth到keras的转换过程中，也用到了onnx。

2，在训练过程中，两框架之间的转换
上述方法只是通过onnx将已经训练好的torch模型转换成tensorflow模型，是否可以在训练过程中，实现两个推理框架之间的相互调用呢，答案是肯定的。
(1) 项目介绍：
在开源项目PyTorch（github项目地址：BlackHC/TfPyTh）中，实现了在模型训练过程中，两个推理框架之间的相互调用。不论是 TensorFlow 还是 PyTorch 计算图，它们都可以包装成一个可微函数，并在另一个框架中高效完成前向与反向传播。很显然，这样的框架交互，能节省很多重写代码的麻烦事。
(2) 转换库 TfPyTh
a，它无需改写已有的代码就能在框架间自由转换。具体而言，TfPyTh 允许我们将 TensorFlow 计算图包装成一个可调用可微分的简单函数，然后 PyTorch 就能直接调用它完成计算。反过来也是同样的，TensorFlow 也能直接调用转换后的 PyTorch 计算图。

b, 因为转换后的模块是可微的，那么正向和反向传播都没什么问题。不过项目作者也表示该项目还不太完美，开源 3 天以来会有一些小的问题。例如张量必须通过 CPU 进行复制与路由，直到 TensorFlow 支持__cuda_array_interface 相关功能才能解决。

目前 TfPyTh 主要支持三大方法：
torch_from_tensorflow：创建一个 PyTorch 可微函数，并给定 TensorFlow 占位符输入计算张量输出；
eager_tensorflow_from_torch：从 PyTorch 创建一个 Eager TensorFlow 函数；
tensorflow_from_torch：从 PyTorch 创建一个 TensorFlow 运算子或张量。
(3) TfPyTh的使用方法
如下所示为 torch_from_tensorflow 的使用案例，我们会用 TensorFlow 创建一个简单的静态计算图，然后传入 PyTorch 张量进行计算。

import tensorflow as tf
import torch as th
import numpy as np
import tfpyth
session = tf.Session()
def get_torch_function():
    a = tf.placeholder(tf.float32, name='a')
    b = tf.placeholder(tf.float32, name='b')
    c = 3 * a + 4 * b * b
    f = tfpyth.torch_from_tensorflow(session, [a, b], c).apply
    return f
f = get_torch_function()
a = th.tensor(1, dtype=th.float32, requires_grad=True)
b = th.tensor(3, dtype=th.float32, requires_grad=True)
x = f(a, b)
assert x == 39.
x.backward()
assert np.allclose((a.grad, b.grad), (3., 24.))

我们可以发现，基本上 TensorFlow 完成的就是一般的运算，例如设置占位符和建立计算流程等。TF 的静态计算图可以通过 session 传递到 TfPyTh 库中，然后就产生了一个新的可微函数。后面我们可以将该函数用于模型的某个计算部分，再进行训练也就没什么问题了。同理，我们可以通过tensorflow_from_torch实现，在tf模型中调用torch的计算模块。

3，将已经训练好的tensorflow模型转化为pytorch模型后，如何减小模型转化带来的偏差。
细节1：input_data的归一化方式，通道顺序(RGB or BGR)是否一致，opencv读入图片是BGR，PIL读入图片是RGB。

细节2：TensorFlow的weight是[kh,kw,in_ch,out_ch]，输入输出[n,h,w,c]，PyTorch的weight是[out_ch,in_ch,kh,kw]，输入输出是[n,c,h,w]。

细节3：batch normlization的moment两个框架的设置正好是相反的，例如tf的moment设定0.99时，PyTorch的moment应该为0.01。

细节4：注意训练好的moving_mean和moving_variance的转化，PyTorch打印named_parameters时是看不到这些的，要写在BN层的state_dict里面，然后PyTorch跑模型的时候记得model.eval()冻结它们。

细节5：padding方式，TensorFlow常用的是“SAME”和"VALID"，会根据图像尺寸，卷积核的size和stride进行计算，保证全覆盖(不会多一行出来没有参与计算)，而PyTorch的padding是每次自己在nn.Conv2d中指定的，并且指定的灵活性还不强，除非自定义conv2d，用nn.zeropadding，反正就是PyTorch在stride=2的层容易出现padding完之后右边和下边出现空白行没有参与计算的情况，这个可以注意查看一下！

总而言之，建议题主先选择同一个数据作为输入分别在TensorFlow和PyTorch上跑一次前向，把每层拿出来进行数据(注意维度顺序不同)的逐点对比，定位问题所在，然后修改掉，估计就能实现TensorFlow和PyTorch结果的完全匹配了！

参考链接：

https://www.zhihu.com/question/387680897
https://blog.csdn.net/learning_tortosie/article/details/109389653
https://www.jianshu.com/p/633c7d1c5cff
https://blog.csdn.net/qq_35447659/article/details/105735504