我们如何将 .pth 模型转换为 .pb 文件? [英] How can we convert a .pth model into .pb file?

问题描述

我已经使用 pytorch 获得了完整的模型，但是我想将 .pth 文件转换为 .pb，可以在 Tensorflow 中使用.有人有什么想法吗?

I have already got the complete model by using pytorch, however I wanna convert the .pth file into .pb, which could be used in Tensorflow. Does anyone have some ideas?

推荐答案

你可以使用 ONNX: Open Neural Network Exchange格式

You can use ONNX: Open Neural Network Exchange Format

将.pth文件转换为.pb 首先，需要将PyTorch中定义的模型导出到ONNX，然后将ONNX模型导入Tensorflow(PyTorch =>ONNX => TensorFlow)

To convert .pth file to .pb First, you need to export a model defined in PyTorch to ONNX and then import the ONNX model into Tensorflow (PyTorch => ONNX => Tensorflow)

这是一个 MNISTModel 到 使用 ONNX 将 PyTorch 模型转换为 Tensorflow 来自 onnx/教程

This is an example of MNISTModel to Convert a PyTorch model to Tensorflow using ONNX from onnx/tutorials

torch.save(model.state_dict(), 'output/mnist.pth')

从文件加载训练好的模型

trained_model = Net()
trained_model.load_state_dict(torch.load('output/mnist.pth'))

# Export the trained model to ONNX
dummy_input = Variable(torch.randn(1, 1, 28, 28)) # one black and white 28 x 28 picture will be the input to the model
torch.onnx.export(trained_model, dummy_input, "output/mnist.onnx")

加载 ONNX 文件

model = onnx.load('output/mnist.onnx')

# Import the ONNX model to Tensorflow
tf_rep = prepare(model)

将 Tensorflow 模型保存到文件中

tf_rep.export_graph('output/mnist.pb')

AS 在评论中由 @tsveti_iko 指出

AS noted by @tsveti_iko in the comment

注意:prepare() 内置在 onnx-tf 中，因此您首先需要像这样通过控制台安装它 pip installonnx-tf，然后像这样在代码中导入它:import onnx from onnx_tf.backend import prepare 然后你就可以按照答案中的描述使用它了.

NOTE: The prepare() is build-in in the onnx-tf, so you first need to install it through the console like this pip install onnx-tf, then import it in the code like this: import onnx from onnx_tf.backend import prepare and after that you can finally use it as described in the answer.

这篇关于我们如何将 .pth 模型转换为 .pb 文件?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！