训练中同时使用tensorflow与pytorch

在开源的 TfPyTh 中,不论是 TensorFlow 还是 PyTorch 计算图,它们都可以包装成一个可微函数,并在另一个框架中高效完成前向与反向传播。

github项目地址:BlackHC/TfPyTh。

神经网络交换格式 ONNX,定义了一种通用的计算图。目前 ONNX 已经原生支持 MXNet、PyTorch 和 Caffe2 等大多数框架,但是像 TensorFlow 或 Keras 之类的只能通过第三方转换器转换为 ONNX 格式。TfPyTh 无需改写已有的代码就能在框架间自由转换。TfPyTh 允许我们将 TensorFlow 计算图包装成一个可调用、可微分的简单函数,然后 PyTorch 就能直接调用它完成计算。反过来也是同样的,TensorFlow 也能直接调用转换后的 PyTorch 计算图。

目前 TfPyTh 主要支持三大方法:

torch_from_tensorflow:创建一个 PyTorch 可微函数,并给定 TensorFlow 占位符输入计算张量输出;

eager_tensorflow_from_torch:从 PyTorch 创建一个 Eager TensorFlow 函数;

tensorflow_from_torch:从 PyTorch 创建一个 TensorFlow 运算子或张量。TfPyTh 示例

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.))

  • 1
    点赞
  • 2
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值