当前位置：服务支持 > 软件文章 > 将PyTorch中的tensor dtype修改与TensorFlow代码改写实践

将PyTorch中的tensor dtype修改与TensorFlow代码改写实践

阅读数 1979

为什么我要折腾模型转换？你也遇到过这个问题

在项目里遇到了个烦心事，之前用PyTorch训练的Query2Title模型，结果要上线的时候发现全得用TensorFlow。这事我挺崩溃的，毕竟用PyTorch练模型已经练出肌肉记忆了，突然要转战TensorFlow，那种感觉就像突然从键盘打字改成手写笔记。

你以为只是文件格式转换？其实藏着不少弯弯绕绕

PyTorch的.pt文件要想变成TensorFlow的.pb文件，中间得经过两个关键步骤。记得第一次尝试的时候，直接对着文档照搬代码，结果报了奇怪的错误。后来才知道，这事儿没简单。

第一步：模型文件的格式迁移

把PyTorch模型导出成ONNX格式，这个过程我一开始觉得挺直接。用torch.onnx.export就能搞定，但后来发现这其实是场硬仗。特别是当我用numpy数组喂数据的时候，一直报维度不匹配的错。

import torchfrom torch.autograd import Variablemodel = YourModel()model.load_state_dict(torch.load('model.pt', map_location='cpu'))dummy_input = Variable(torch.randint(0, 10000, (1, 20)))torch.onnx.export(model, dummy_input, 'model.onnx')

这段代码其实也有陷阱，在2026年现在，onnx版本和PyTorch版本得严格匹配。我试过用PyTorch 1.9导出的模型，结果在ONNX转换的时候报错，搞了大半天才发现版本不兼容的问题。

转换过程中的那些坑

当到第二步转换成TensorFlow模型时，情况更复杂。我一开始照搬官方文档写的代码，结果就是报错：

model_onnx = onnx.load('model.onnx')tf_rep = prepare(model_onnx)tf_rep.export_graph('model.pb')

这代码在2026年已经不适用了，onnx-tf的兼容性太差。我改成命令行方式才搞定，这种经验真的值得记录。

版本适配是个大问题

我最早用TensorFlow 1.12版本的时候，转换出来的pb文件根本没法跑。结果发现这是onnx-tf版本的问题，1.12只能用onnx-tf 1.3。后来升级到1.13之后，问题就解决了。

| Tensorflow版本 | 可用onnx-tf版本 |

|----------------|----------------|

| 1.12 | 1.3 |

| 1.13 | 1.7 |

| 2.0+ | 使用最新版 |

运行环境更要小心

最离谱的是GPU环境和CPU环境的差异。某天我测试的时候发现，用GPU跑出来的结果和CPU的结果不一样。究其原因，发现max_pool层在CPU环境下只支持NHWC格式，而PyTorch默认是NCHW。

os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

这段代码就是解决之道。虽然没有GPU环境的话根本不用管，但预防性写上总没错。我现在测试模型都在CPU上跑，报错概率能降低不少。

实操时候的那些小招数

其实我测试的时候发现个有意思的事。用PyTorch模型跑的时候，输入维度是1x20，结果TensorFlow导出的模型输入维度变成了20。这个差异差点让我误以为模型训练出错了。

def load_pb(path_to_pb):with tf.gfile.GFile(path_to_pb, 'rb') as f:graph_def = tf.GraphDef()graph_def.ParseFromString(f.read())with tf.Graph().as_default() as graph:tf.import_graph_def(graph_def, name='')return graph

这代码前提是得知道输入输出节点的名称。试过用'inputs:0'当输入节点，结果发现输出节点名是'div_3:0'，这个需要自己去计算图里找。

我的测试案例记录

用这个模型做了个实际测试，导入两个查询样本。输入维度是1x20，输出是16维的向量。用numpy的dot计算相似度，结果发现差异居然达到了0.3。这个数据挺有说服力的。

query1 = np.array([[4158, 7811, 6653, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=np.int64)query2 = np.array([[9914, 10859, 6907, 8719, 7098, 8861, 4158, 10785, 6299, 1264, 1612, 10285, 6973, 7811, 0, 0]], dtype=np.int64)

这两个样本我折磨了整整三天，终于在2026年3月找到了稳定的迁移方案。说实话，到现在我还在用PyTorch做实验，但每次部署都得折腾这些转换步骤。