Как мы можем преобразовать модель .pth в файл .pb?

Question

Я уже получил полную модель с помощью pytorch, однако я хочу преобразовать файл .pth в .pb, который можно использовать в Tensorflow. У кого-нибудь есть идеи?

Answer 1

Вы можете использовать ONNX : открыть формат обмена нейронной сетью

Чтобы преобразовать файл .pth в .pb Сначала необходимо экспортировать модель, определенную в PyTorch, в ONNX, а затемимпортировать модель ONNX в Tensorflow (PyTorch => ONNX => Tensorflow)

Это пример MNISTModel в Преобразование модели PyTorch в Tensorflow с использованием ONNX из 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')