Comment convertir un modèle .pth en fichier .pb ?

Question

J'ai déjà obtenu le modèle complet en utilisant pytorch, mais je veux convertir le fichier .pth en .pb, qui pourrait être utilisé dans Tensorflow. Quelqu'un a-t-il une idée ?

Answer 1

Vous pouvez utiliser ONNX : Format d'échange ouvert de réseaux neuronaux

Pour convertir .pth pour .pb Tout d'abord, vous devez exporter un modèle défini dans PyTorch vers ONNX, puis importer le modèle ONNX dans Tensorflow (PyTorch => ONNX => Tensorflow).

Voici un exemple de MNISTModel pour Convertir un modèle PyTorch vers Tensorflow en utilisant ONNX de onnx/tutoriels

Sauvegarder le modèle formé dans un fichier

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

Charger le modèle formé à partir du fichier

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

Charger le fichier ONNX

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

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

Enregistrer le modèle Tensorflow dans un fichier

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

Comme l'a noté @tsveti_iko dans le commentaire

NOTE : Le prepare() est intégré dans le onnx-tf Vous devez donc d'abord l'installer via la console comme suit pip install onnx-tf puis l'importer dans le code comme ceci : import onnx from onnx_tf.backend import prepare et après cela, vous pouvez enfin l'utiliser comme décrit dans la réponse.