Pytorch 和Tensorflow之间的相互转化

转载于：https://mp.weixin.qq.com/s/LzbJ_QuRqWJOj_ZRZ5P66A

1. 介绍

• 在早两天开源的 TfPyTh 中，不论是 TensorFlow 还是 PyTorch 计算图，它们都可以包装成一个可微函数，并在另一个框架中高效完成前向与反向传播。很显然，这样的框架交互，能节省很多重写代码的麻烦事。
• 项目地址：https://github.com/BlackHC/TfPyTh

2.神奇的转换库 TfPyTh

• 既然 ONNX 无法解决训练问题，那么就轮到 TfPyTh 这类项目出场了，它无需改写已有的代码就能在框架间自由转换。具体而言，TfPyTh 允许我们将 TensorFlow 计算图包装成一个可调用可微分的简单函数，然后 PyTorch 就能直接调用它完成计算。反过来也是同样的，TensorFlow 也能直接调用转换后的 PyTorch 计算图。
• 因为转换后的模块是可微的，那么正向和反向传播都没什么问题。不过项目作者也表示该项目还不太完美，开源 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 库中，然后就产生了一个新的可微函数。后面我们可以将该函数用于模型的某个计算部分，再进行训练也就没什么问题了.