Converting Tensorflow model (weights) to pytorch

Hi,

Looking for ways to convert a custom tensorflow trained model to pytorch equivalent that can be fine-tuned further in pytorch.

Thanks

I don’t know if there are tools to convert the TF model automatically to PyTorch and think you would have to rewrite it manually in PyTorch.
A definition of a custom model can be found in this tutorial and might be a good starter.

Once the model architecture is created in PyTorch, you could convert the pretrained weights from TF to PyTorch. @tom and I did the same for the StyleGAN model in this notebook so you could take a look at the implementation.

4 Likes

Hi @ptrblck, thanks that helped. I was able extract the sequence of layer information and corresponding weights using tf.all_variables().

The following are the components I am trying to transfer weights from tensorflow to pytorch now,

TF layers and weights:

  1. Batchnorm layer: BatchNorm/beta:0, BatchNorm/moving_mean:0 and moving_variance:0
  2. Conv layer: Conv/weights:0 and Conv/biases:0
  3. Conv2d_transpose layer: Conv2d_transpose_1/weights:0 and Conv2d_transpose_1/biases:0

Can I just create the same sequence of layers in pytorch and assign these? Will there be computational flow differences between pytorch and tensorflow?
Wondering if you have idea on conversion for those particular layers.

The notebook also shows how to transfer these parameter to PyTorch.
In particular, check out weight_translate, which permutes the conv filters etc. to match the expected shape for PyTorch layers.

1 Like