Leonardo-Blanger / detr_tensorflow

A Tensorflow implementation of the DETR object detection architecture.
MIT License
62 stars 15 forks source link

How to Convert the Pytorch Weights to H5 #2

Closed ktobah closed 3 years ago

ktobah commented 4 years ago

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.

ktobah commented 3 years ago

Well, I figured how to convert the model from Pytorch to Tensorflow SavedModel.

If someone needs it.

  1. Convert pytorch model to onnx format: load your pytorch detr model and use .eval() to convert it to an inference model. Then create a dummy variable input and feed it to the export() method. Make sure to set 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)
  1. Convert onnx to tf saved_model (for tf serving):
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.

ktobah commented 3 years ago

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.