Understanding difference between TensorFlow and PyTorch implementations of Unet

[simondgreenhill]

I am trying to convert the weights from a segmentation-models-based model that was trained in TensorFlow to PyTorch. I am using a U-net with an efficientnet-b0 backbone.

I noticed a mismatch in the dimensions of the parameters across the PyTorch and and TensorFlow implementations in the convolutional layers of the decoder. What is the source of this discrepancy, and how can I make the architectures match exactly?

Reproducible example:

# set up models

# pytorch
import torch.nn
import segementation_models_pytorch as smp
from functools import partial

pytorch_model = smp.Unet('efficientnet-b0', classes=4, activation=partial(nn.Softmax, dim=0))


# tensorflow
import segmentation_models as sm
tf_model = sm.Unet('efficientnetb0', classes=4, activation='softmax', encoder_freeze=False)

Observe the mismatch in the convolutional layers of the decoder, taking the first one as an example:

layer = tf_test.get_layer('decoder_stage0a_conv')
print(layer.weights[0].shape)
# TensorShape([3, 3, 1952, 256])
print(t.decoder.blocks[0].conv1[0].weight.shape)
# torch.Size([16, 432, 3, 3])

Note how the TensorFlow version has 1952 in channels, and the PyTorch version has 432 in-channels. Why is this the case?

[pinnintipraneethkumar]

Hii,
Try to convert the tensorflow model to ONNX and from ONNX to pytorch model, hope it will resolve this problem.

[github-actions]

This issue is stale because it has been open 60 days with no activity. Remove stale label or comment or this will be closed in 7 days.

[github-actions]

This issue was closed because it has been stalled for 7 days with no activity.

[hariouat]

I am trying to convert the weights from a segmentation-models-based model that was trained in TensorFlow to PyTorch. I am using a U-net with an efficientnet-b0 backbone.

I noticed a mismatch in the dimensions of the parameters across the PyTorch and and TensorFlow implementations in the convolutional layers of the decoder. What is the source of this discrepancy, and how can I make the architectures match exactly?

Reproducible example:

# set up models

# pytorch
import torch.nn
import segementation_models_pytorch as smp
from functools import partial

pytorch_model = smp.Unet('efficientnet-b0', classes=4, activation=partial(nn.Softmax, dim=0))


# tensorflow
import segmentation_models as sm
tf_model = sm.Unet('efficientnetb0', classes=4, activation='softmax', encoder_freeze=False)

Observe the mismatch in the convolutional layers of the decoder, taking the first one as an example:

layer = tf_test.get_layer('decoder_stage0a_conv')
print(layer.weights[0].shape)
# TensorShape([3, 3, 1952, 256])
print(t.decoder.blocks[0].conv1[0].weight.shape)
# torch.Size([16, 432, 3, 3])

Note how the TensorFlow version has 1952 in channels, and the PyTorch version has 432 in-channels. Why is this the case?

Hello,
I m facing the same problem.
Do you have any solution pelase?
Thanks in advance