test_models.py

import pytest
import torch

import segmentation_models_pytorch as smp  # noqa


def get_encoders():
    exclude_encoders = [
        "senet154",
        "resnext101_32x16d",
        "resnext101_32x32d",
        "resnext101_32x48d",
    ]
    encoders = smp.encoders.get_encoder_names()
    encoders = [e for e in encoders if e not in exclude_encoders]
    encoders.append("tu-resnet34")  # for timm universal encoder
    return encoders


ENCODERS = get_encoders()
DEFAULT_ENCODER = "resnet18"


def get_sample(model_class):
    if model_class in [
        smp.FPN,
        smp.Linknet,
        smp.Unet,
        smp.UnetPlusPlus,
        smp.MAnet,
        smp.Segformer,
    ]:
        sample = torch.ones([1, 3, 64, 64])
    elif model_class in [smp.PAN, smp.DeepLabV3, smp.DeepLabV3Plus]:
        sample = torch.ones([2, 3, 128, 128])
    elif model_class in [smp.PSPNet, smp.UPerNet]:
        # Batch size 2 needed due to nn.BatchNorm2d not supporting (1, C, 1, 1) input
        # from PSPModule pooling in PSPNet/UPerNet.
        sample = torch.ones([2, 3, 64, 64])
    else:
        raise ValueError("Not supported model class {}".format(model_class))
    return sample


def _test_forward(model, sample, test_shape=False):
    with torch.no_grad():
        out = model(sample)
    if test_shape:
        assert out.shape[2:] == sample.shape[2:]


def _test_forward_backward(model, sample, test_shape=False):
    out = model(sample)
    out.mean().backward()
    if test_shape:
        assert out.shape[2:] == sample.shape[2:]


@pytest.mark.parametrize("encoder_name", ENCODERS)
@pytest.mark.parametrize("encoder_depth", [3, 5])
@pytest.mark.parametrize(
    "model_class",
    [
        smp.FPN,
        smp.PSPNet,
        smp.Linknet,
        smp.Unet,
        smp.UnetPlusPlus,
        smp.MAnet,
        smp.UPerNet,
        smp.Segformer,
    ],
)
def test_forward(model_class, encoder_name, encoder_depth, **kwargs):
    if (
        model_class is smp.Unet
        or model_class is smp.UnetPlusPlus
        or model_class is smp.MAnet
    ):
        kwargs["decoder_channels"] = (16, 16, 16, 16, 16)[-encoder_depth:]
    if model_class in [smp.UnetPlusPlus, smp.Linknet] and encoder_name.startswith(
        "mit_b"
    ):
        return  # skip mit_b*
    if (
        model_class is smp.FPN
        and encoder_name.startswith("mit_b")
        and encoder_depth != 5
    ):
        return  # skip mit_b*
    model = model_class(
        encoder_name, encoder_depth=encoder_depth, encoder_weights=None, **kwargs
    )
    sample = get_sample(model_class)
    model.eval()
    if encoder_depth == 5 and model_class != smp.PSPNet:
        test_shape = True
    else:
        test_shape = False

    _test_forward(model, sample, test_shape)


@pytest.mark.parametrize(
    "model_class",
    [
        smp.PAN,
        smp.FPN,
        smp.PSPNet,
        smp.Linknet,
        smp.Unet,
        smp.UnetPlusPlus,
        smp.MAnet,
        smp.DeepLabV3,
        smp.DeepLabV3Plus,
        smp.UPerNet,
        smp.Segformer,
    ],
)
def test_forward_backward(model_class):
    sample = get_sample(model_class)
    model = model_class(DEFAULT_ENCODER, encoder_weights=None)
    _test_forward_backward(model, sample)


@pytest.mark.parametrize(
    "model_class",
    [
        smp.PAN,
        smp.FPN,
        smp.PSPNet,
        smp.Linknet,
        smp.Unet,
        smp.UnetPlusPlus,
        smp.MAnet,
        smp.DeepLabV3,
        smp.DeepLabV3Plus,
        smp.UPerNet,
        smp.Segformer,
    ],
)
def test_aux_output(model_class):
    model = model_class(
        DEFAULT_ENCODER, encoder_weights=None, aux_params=dict(classes=2)
    )
    sample = get_sample(model_class)
    label_size = (sample.shape[0], 2)
    mask, label = model(sample)
    assert label.size() == label_size


@pytest.mark.parametrize("upsampling", [2, 4, 8])
@pytest.mark.parametrize("model_class", [smp.FPN, smp.PSPNet])
def test_upsample(model_class, upsampling):
    default_upsampling = 4 if model_class is smp.FPN else 8
    model = model_class(DEFAULT_ENCODER, encoder_weights=None, upsampling=upsampling)
    sample = get_sample(model_class)
    mask = model(sample)
    assert mask.size()[-1] / 64 == upsampling / default_upsampling


@pytest.mark.parametrize("model_class", [smp.FPN])
@pytest.mark.parametrize("in_channels", [1, 2, 4])
def test_in_channels(model_class, in_channels):
    sample = torch.ones([1, in_channels, 64, 64])
    model = model_class(DEFAULT_ENCODER, encoder_weights=None, in_channels=in_channels)
    model.eval()
    with torch.no_grad():
        model(sample)

    assert model.encoder._in_channels == in_channels


@pytest.mark.parametrize("encoder_name", ENCODERS)
def test_dilation(encoder_name):
    if (
        encoder_name in ["inceptionresnetv2", "xception", "inceptionv4"]
        or encoder_name.startswith("vgg")
        or encoder_name.startswith("densenet")
        or encoder_name.startswith("timm-res")
        or encoder_name.startswith("mit_b")
    ):
        return

    encoder = smp.encoders.get_encoder(encoder_name, output_stride=16)

    encoder.eval()
    with torch.no_grad():
        sample = torch.ones([1, 3, 64, 64])
        output = encoder(sample)

    shapes = [out.shape[-1] for out in output]
    assert shapes == [64, 32, 16, 8, 4, 4]  # last downsampling replaced with dilation


if __name__ == "__main__":
    pytest.main([__file__])