Convert PyTorch model to ONNX to TF 2 SavedModel

pt_to_onnx_to_pb.py

# Use torch.onnx.export function

model_pytorch = SimpleModel(input_size=input_size, hidden_sizes=hidden_sizes, output_size=output_size)
model_pytorch.load_state_dict(torch.load('./models/model_simple.pt'))

sample_input = torch.from_numpy(X_test[0].reshape(1, -1)).float().to(device)
sample_output = model_pytorch(dummy_input)

# Export to ONNX format
torch.onnx.export(model_pytorch, sample_input, './models/model_simple.onnx', input_names=['test_input'], output_names=['test_output'])

# Load ONNX model and convert to TensorFlow format
model_onnx = onnx.load('./models/model_simple.onnx')

tf_rep = prepare(model_onnx)

# Export model as .pb file
tf_rep.export_graph('./models/model_simple.pb')