Closed ktobah closed 3 years ago
Well, I figured how to convert the model from Pytorch to Tensorflow SavedModel.
If someone needs it.
opset_version=11
dummy_input = Variable(torch.randn(1, 3, 800, 1777))
# you may change the names
input_names = ['input_image']
output_names = ['pred_logits', 'pred_boxes']
# I use dynamic axes to make the model accept a variable batch size, width, and height for the input images
dynamic_axes = {'input_image': {0: 'batch_size', 2: 'width', 3: 'height'}}
torch.onnx.export(model, dummy_input, "model.onnx", input_names=input_names, output_names=output_names,
dynamic_axes=dynamic_axes, verbose=1, opset_version=11)
import onnx
from onnx_tf.backend import prepare
model = onnx.load('model.onnx')
# Import the ONNX model to Tensorflow
tf_rep = prepare(model)
# don't use the ".pb" in the name of the exported file, so that it creates a proper folder for the weights
tf_rep.export_graph("detr")
I am closing this.
I wanted to mention that the input preprocessing in this repo does not give the same result as pytorch's preprocessing. So, I keep using pytorch preprocessing.
Thank you for this work.
I have a question rather than an issue. I have trained a DETR model using the PyTorch implementation, and now I would like to convert the model into a TF model. I think your implementation allows this, except that I am not sure how you converted the weights to a .h5 file.
Would you please shed some light on this? Thank you.