问题描述
我已经使用pytorch训练了图像分类模型。
现在，我想将其从研究转移到生产流程。
我正在考虑使用TensorFlow扩展。我非常怀疑是否可以在TensorFlow扩展管道中使用经过PyTorch训练的模型（我可以将训练后的模型转换为ONNX，然后转换为Tensorflow兼容格式）。
我不想将训练部分重写并重新训练到TensorFlow，因为这将是巨大的开销。
有可能还是有更好的方法来生产PyTorch训练过的模型？
解决方法
只要您使用标准层，您就应该能够使用 ONNX 将您的 PyTorch 图像分类模型转换为 Tensorflow 格式。我建议进行转换，然后查看两个模型摘要以确保它们相对相似。此外，进行一些测试以确保您转换后的模型能够处理您拥有的任何特定边缘情况。确认转换后的模型有效后，将模型保存为 TF SavedModel 格式，然后您应该可以在 Tensorflow Extended (TFX) 中使用它。
有关转换过程的更多信息，请参阅本教程：https://learnopencv.com/pytorch-to-tensorflow-model-conversion/