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?

This is a cross-post of the issue #840 in the segmentation_models.pytorch repo.