Support for ConvNeXt backbone / timm v 0.5.4

[jkuechler-ce]

PyTorch image models supports ConvNeXt since version 0.5.4, which could be an interesting backbone for segmentation. Right now segmentation_models.pytorch is still using timm==0.4.12. Would it be possible to switch to the newer version?

[qubvel]

Hi, you could try to update timm after the smp installation and check if it still works.

[qubvel]

I will definitely update it later

[jkuechler-ce]

Hi,

upgrading timm after the smp installation actually seems to work, thanks for the idea!

Unfortunately creating a UNet model with a ConvNeXt backbone gives an output that has half the width and height of the input:

import segmentation_models_pytorch as smp
import timm
import torch

assert timm.__version__ == '0.5.4'

# need to use encoder_depth=4, because convnext_tiny isn't that large
model = smp.Unet('tu-convnext_tiny', classes=11, activation='softmax2d', encoder_depth=4, decoder_channels=(128,64,32,16))

dummy_input = torch.rand(1, 3, 224, 224)
output = model(dummy_input)

print(output.shape)
# gives: torch.Size([1, 11, 112, 112])

Do you have an idea what I'm doing wrong here?

[qubvel]

Does it have the same behaviour with encoder_depth=5?

[jkuechler-ce]

With encoder_depth=5 there is an error, because the backbone only has 4 feature outputs:

>>> model = smp.Unet('tu-convnext_tiny', classes=11, activation='softmax2d', encoder_depth=5)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File .\projects\convnext-pytorch\env\lib\site-packages\segmentation_models_pytorch\unet\model.py", line 65, in __init__
    self.encoder = get_encoder(
  File .\projects\convnext-pytorch\env\lib\site-packages\segmentation_models_pytorch\encoders\__init__.py", line 50, in get_encoder
    encoder = TimmUniversalEncoder(
  File .\projects\convnext-pytorch\env\lib\site-packages\segmentation_models_pytorch\encoders\timm_universal.py", line 21, in __init__
    self.model = timm.create_model(name, **kwargs)
  File .\projects\convnext-pytorch\env\lib\site-packages\timm\models\factory.py", line 74, in create_model
    model = create_fn(pretrained=pretrained, **kwargs)
  File .\projects\convnext-pytorch\env\lib\site-packages\timm\models\convnext.py", line 332, in convnext_tiny
    model = _create_convnext('convnext_tiny', pretrained=pretrained, **model_args)
  File .\projects\convnext-pytorch\env\lib\site-packages\timm\models\convnext.py", line 320, in _create_convnext
    model = build_model_with_cfg(
  File .\projects\convnext-pytorch\env\lib\site-packages\timm\models\helpers.py", line 485, in build_model_with_cfg
    model = feature_cls(model, **feature_cfg)
  File .\projects\convnext-pytorch\env\lib\site-packages\timm\models\features.py", line 227, in __init__
    super(FeatureListNet, self).__init__(
  File .\projects\convnext-pytorch\env\lib\site-packages\timm\models\features.py", line 184, in __init__
    return_layers = _get_return_layers(self.feature_info, out_map)
  File .\projects\convnext-pytorch\env\lib\site-packages\timm\models\features.py", line 146, in _get_return_layers
    module_names = feature_info.module_name()
  File .\projects\convnext-pytorch\env\lib\site-packages\timm\models\features.py", line 75, in module_name
    return self.get('module', idx)
  File .\projects\convnext-pytorch\env\lib\site-packages\timm\models\features.py", line 43, in get
    return [self.info[i][key] for i in self.out_indices]
  File .\projects\convnext-pytorch\env\lib\site-packages\timm\models\features.py", line 43, in <listcomp>
    return [self.info[i][key] for i in self.out_indices]
IndexError: list index out of range

[jkuechler-ce]

The problem seems to be rooted in the ConvNeXt architecture. The backbone starts with a convolution with a kernel size of 4 and a stride of 4, which is not suited for UNet (which is upsampling by a factor of 2 only).

[ynhuhu]

Unknown model (convnext_tiny)
Convnext is not available.Can convnext be added to SMP?

[ynhuhu]

Sorry.I am wrong.It is available after timm is updated.

[carbocation]

It seems that one solution would be to make the strides configurable, so that the unusual stride choice in the first convolution of ConvNext could be handled.