@@ -4,8 +4,8 @@ from super_gradients.common.aws_connection import AWSConnector
 from super_gradients.common.data_connection import S3Connector
 from super_gradients.common.data_interface import DatasetDataInterface, ADNNModelRepositoryDataInterfaces
 from super_gradients.common.environment.env_helpers import init_trainer, is_distributed
-from super_gradients.common.data_types import StrictLoad, DeepLearningTask, EvaluationType, MultiGPUMode
+from super_gradients.common.data_types import StrictLoad, DeepLearningTask, EvaluationType, MultiGPUMode, UpsampleMode
 
 __all__ = ['explicit_params_validation', 'singleton', 'AWSConnector', 'DatasetDataInterface',
            'ADNNModelRepositoryDataInterfaces', 'S3Connector', 'init_trainer', 'is_distributed',
-           'StrictLoad', 'DeepLearningTask', 'EvaluationType', 'MultiGPUMode']
+           'StrictLoad', 'DeepLearningTask', 'EvaluationType', 'MultiGPUMode', 'UpsampleMode']

@@ -1,3 +1,4 @@
-from super_gradients.common.data_types.enum import StrictLoad, DeepLearningTask, MultiGPUMode, EvaluationType
+from super_gradients.common.data_types.enum import StrictLoad, DeepLearningTask, MultiGPUMode, EvaluationType,\
+    UpsampleMode
 
-__all__ = ['StrictLoad', 'DeepLearningTask', 'EvaluationType', 'MultiGPUMode']
+__all__ = ['StrictLoad', 'DeepLearningTask', 'EvaluationType', 'MultiGPUMode', 'UpsampleMode']

@@ -2,6 +2,7 @@ from super_gradients.common.data_types.enum.strict_load import StrictLoad
 from super_gradients.common.data_types.enum.deep_learning_task import DeepLearningTask
 from super_gradients.common.data_types.enum.evaluation_type import EvaluationType
 from super_gradients.common.data_types.enum.multi_gpu_mode import MultiGPUMode
+from super_gradients.common.data_types.enum.upsample_mode import UpsampleMode
 
 
-__all__ = ['StrictLoad', 'DeepLearningTask', 'EvaluationType', 'MultiGPUMode']
+__all__ = ['StrictLoad', 'DeepLearningTask', 'EvaluationType', 'MultiGPUMode', 'UpsampleMode']

from enum import Enum


class UpsampleMode(Enum):
    NEAREST = "nearest"
    BILINEAR = "bilinear"
    BICUBIC = "bicubic"
    SNPE_BILINEAR = "snpe_bilinear"

@@ -21,6 +21,7 @@ from super_gradients.training.models.segmentation_models.stdc import STDC1Classi
 
 from super_gradients.training.models.kd_modules.kd_module import KDModule
 from super_gradients.training.models.classification_models.beit import beit_base_patch16_224, beit_large_patch16_224
+from super_gradients.training.models.segmentation_models.ppliteseg import PPLiteSegT, PPLiteSegB
 # IF YOU UPDATE THE ARCHITECTURE DICT PLEASE ALSO UPDATE THE ENUM CLASS DOWN BELOW.
 
 
@@ -112,7 +113,13 @@ ARCHITECTURES = {"resnet18": resnet.ResNet18,
                  "vit_large": vit_large,
                  "vit_huge": vit_huge,
                  "beit_base_patch16_224": beit_base_patch16_224,
-                 "beit_large_patch16_224": beit_large_patch16_224
+                 "beit_large_patch16_224": beit_large_patch16_224,
+                 "pp_lite_t_seg": PPLiteSegT,
+                 "pp_lite_t_seg50": PPLiteSegT,
+                 "pp_lite_t_seg75": PPLiteSegT,
+                 "pp_lite_b_seg": PPLiteSegB,
+                 "pp_lite_b_seg50": PPLiteSegB,
+                 "pp_lite_b_seg75": PPLiteSegB,
                  }
 
 KD_ARCHITECTURES = {

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Union, List, Tuple

from super_gradients.training.utils.module_utils import ConvBNReLU, make_upsample_module
from super_gradients.common import UpsampleMode
from super_gradients.training.models.segmentation_models.stdc import SegmentationHead, AbstractSTDCBackbone,\
    STDC1Backbone, STDC2Backbone
from super_gradients.training.models.segmentation_models.segmentation_module import SegmentationModule
from super_gradients.training.utils import HpmStruct, get_param


class SPPM(nn.Module):
    """
    Simple Pyramid Pooling context Module.
    """
    def __init__(self,
                 in_channels: int,
                 inter_channels: int,
                 out_channels: int,
                 pool_sizes: List[Union[int, Tuple[int, int]]],
                 upsample_mode: Union[UpsampleMode, str] = UpsampleMode.BILINEAR,
                 align_corners: bool = False):
        """
        :param inter_channels: num channels in each pooling branch.
        :param out_channels: The number of output channels after pyramid pooling module.
        :param pool_sizes: spatial output sizes of the pooled feature maps.
        """
        super().__init__()
        self.branches = nn.ModuleList([
            nn.Sequential(
                nn.AdaptiveAvgPool2d(pool_size),
                ConvBNReLU(in_channels, inter_channels, kernel_size=1, bias=False),
            ) for pool_size in pool_sizes
        ])
        self.conv_out = ConvBNReLU(inter_channels, out_channels, kernel_size=3, padding=1, bias=False)
        self.out_channels = out_channels
        self.upsample_mode = upsample_mode
        self.align_corners = align_corners
        self.pool_sizes = pool_sizes

    def forward(self, x):
        out = None
        input_shape = x.shape[2:]
        for branch in self.branches:
            y = branch(x)
            y = F.interpolate(y, size=input_shape, mode=self.upsample_mode, align_corners=self.align_corners)
            out = y if out is None else out + y
        out = self.conv_out(out)
        return out

    def prep_model_for_conversion(self, input_size: Union[tuple, list], stride_ratio: int = 32, **kwargs):
        """
        Replace Global average pooling with fixed kernels Average pooling, since dynamic kernel sizes are not supported
        when compiling to ONNX: `Unsupported: ONNX export of operator adaptive_avg_pool2d, input size not accessible.`
        """
        input_size = [x / stride_ratio for x in input_size[-2:]]
        for branch in self.branches:
            global_pool: nn.AdaptiveAvgPool2d = branch[0]
            out_size = global_pool.output_size
            out_size = out_size if isinstance(out_size, (tuple, list)) else (out_size, out_size)
            kernel_size = [int(i / o) for i, o in zip(input_size, out_size)]
            branch[0] = nn.AvgPool2d(kernel_size=kernel_size, stride=kernel_size)


class UAFM(nn.Module):
    """
    Unified Attention Fusion Module, which uses mean and max values across the spatial dimensions.
    """
    def __init__(self,
                 in_channels: int,
                 skip_channels: int,
                 out_channels: int,
                 up_factor: int,
                 upsample_mode: Union[UpsampleMode, str] = UpsampleMode.BILINEAR,
                 align_corners: bool = False):
        """
        :params in_channels: num_channels of input feature map.
        :param skip_channels: num_channels of skip connection feature map.
        :param out_channels: num out channels after features fusion.
        :param up_factor: upsample scale factor of the input feature map.
        :param upsample_mode: see UpsampleMode for valid options.
        """
        super().__init__()
        self.conv_atten = nn.Sequential(
            ConvBNReLU(4, 2, kernel_size=3, padding=1, bias=False),
            ConvBNReLU(2, 1, kernel_size=3, padding=1, bias=False, use_activation=False)
        )

        self.proj_skip = nn.Identity() if skip_channels == in_channels else \
            ConvBNReLU(skip_channels, in_channels, kernel_size=3, padding=1, bias=False)
        self.up_x = nn.Identity() if up_factor == 1 else \
            make_upsample_module(scale_factor=up_factor, upsample_mode=upsample_mode, align_corners=align_corners)
        self.conv_out = ConvBNReLU(in_channels, out_channels, kernel_size=3, padding=1, bias=False)

    def forward(self, x, skip):
        """
        :param x: input feature map to upsample before fusion.
        :param skip: skip connection feature map.
        """
        x = self.up_x(x)
        skip = self.proj_skip(skip)

        atten = torch.cat([
            *self._avg_max_spatial_reduce(x, use_concat=False),
            *self._avg_max_spatial_reduce(skip, use_concat=False)
        ], dim=1)
        atten = self.conv_atten(atten)
        atten = torch.sigmoid(atten)

        out = x * atten + skip * (1 - atten)
        out = self.conv_out(out)
        return out

    @staticmethod
    def _avg_max_spatial_reduce(x, use_concat: bool = False):
        reduced = [
            torch.mean(x, dim=1, keepdim=True),
            torch.max(x, dim=1, keepdim=True)[0]
        ]
        if use_concat:
            reduced = torch.cat(reduced, dim=1)
        return reduced


class PPLiteSegEncoder(nn.Module):
    """
    Encoder for PPLiteSeg, include backbone followed by a context module.
    """
    def __init__(self,
                 backbone: AbstractSTDCBackbone,
                 projection_channels_list: List[int],
                 context_module: nn.Module):
        super().__init__()
        self.backbone = backbone
        self.context_module = context_module
        feats_channels = backbone.get_backbone_output_number_of_channels()
        self.proj_convs = nn.ModuleList([
            ConvBNReLU(feat_ch, proj_ch, kernel_size=3, padding=1, bias=False)
            for feat_ch, proj_ch in zip(feats_channels, projection_channels_list)
        ])
        self.projection_channels_list = projection_channels_list

    def get_output_number_of_channels(self) -> List[int]:
        channels_list = self.projection_channels_list
        if hasattr(self.context_module, "out_channels"):
            channels_list.append(self.context_module.out_channels)
        return channels_list

    def forward(self, x):
        feats = self.backbone(x)
        y = self.context_module(feats[-1])
        feats = [conv(f) for conv, f in zip(self.proj_convs, feats)]
        return feats + [y]


class PPLiteSegDecoder(nn.Module):
    """
    PPLiteSegDecoder using UAFM blocks to fuse feature maps.
    """
    def __init__(self,
                 encoder_channels: List[int],
                 up_factors: List[int],
                 out_channels: List[int],
                 upsample_mode,
                 align_corners: bool):
        super().__init__()
        # Make a copy of channels list, to prevent out of scope changes.
        encoder_channels = encoder_channels.copy()
        encoder_channels.reverse()
        in_channels = encoder_channels.pop(0)

        # TODO - assert argument length
        self.up_stages = nn.ModuleList()
        for skip_ch, up_factor, out_ch in zip(encoder_channels, up_factors, out_channels):
            self.up_stages.append(UAFM(
                in_channels=in_channels,
                skip_channels=skip_ch,
                out_channels=out_ch,
                up_factor=up_factor,
                upsample_mode=upsample_mode,
                align_corners=align_corners
            ))
            in_channels = out_ch

    def forward(self, feats: List[torch.Tensor]):
        feats.reverse()
        x = feats.pop(0)
        for up_stage, skip in zip(self.up_stages, feats):
            x = up_stage(x, skip)
        return x


class PPLiteSegBase(SegmentationModule):
    """
    The PP_LiteSeg implementation based on PaddlePaddle.
    The original article refers to "Juncai Peng, Yi Liu, Shiyu Tang, Yuying Hao, Lutao Chu,
    Guowei Chen, Zewu Wu, Zeyu Chen, Zhiliang Yu, Yuning Du, Qingqing Dang,Baohua Lai,
    Qiwen Liu, Xiaoguang Hu, Dianhai Yu, Yanjun Ma. PP-LiteSeg: A Superior Real-Time Semantic
    Segmentation Model. https://arxiv.org/abs/2204.02681".
    """
    def __init__(self,
                 num_classes,
                 backbone: AbstractSTDCBackbone,
                 projection_channels_list: List[int],
                 sppm_inter_channels: int,
                 sppm_out_channels: int,
                 sppm_pool_sizes: List[int],
                 sppm_upsample_mode: Union[UpsampleMode, str],
                 align_corners: bool,
                 decoder_up_factors: List[int],
                 decoder_channels: List[int],
                 decoder_upsample_mode: Union[UpsampleMode, str],
                 head_scale_factor: int,
                 head_upsample_mode: Union[UpsampleMode, str],
                 head_mid_channels: int,
                 dropout: float,
                 use_aux_heads: bool,
                 aux_hidden_channels: List[int],
                 aux_scale_factors: List[int]
                 ):
        """
        :param backbone: Backbone nn.Module should implement the abstract class `AbstractSTDCBackbone`.
        :param projection_channels_list: channels list to project encoder features before fusing with the decoder
            stream.
        :param sppm_inter_channels: num channels in each sppm pooling branch.
        :param sppm_out_channels: The number of output channels after sppm module.
        :param sppm_pool_sizes: spatial output sizes of the pooled feature maps.
        :param sppm_upsample_mode: Upsample mode to original size after pooling.
        :param decoder_up_factors: list upsample factor per decoder stage.
        :param decoder_channels: list of num_channels per decoder stage.
        :param decoder_upsample_mode: upsample mode in decoder stages, see UpsampleMode for valid options.
        :param head_scale_factor: scale factor for final the segmentation head logits.
        :param head_upsample_mode: upsample mode to final prediction sizes, see UpsampleMode for valid options.
        :param head_mid_channels: num of hidden channels in segmentation head.
        :param use_aux_heads: set True when training, output extra Auxiliary feature maps from the encoder module.
        :param aux_hidden_channels: List of hidden channels in auxiliary segmentation heads.
        :param aux_scale_factors: list of uppsample factors for final auxiliary heads logits.
        """
        super().__init__(use_aux_heads=use_aux_heads)

        # Init Encoder
        backbone_out_channels = backbone.get_backbone_output_number_of_channels()
        assert len(backbone_out_channels) == len(projection_channels_list), \
            f"The length of backbone outputs ({backbone_out_channels}) should match the length of projection channels" \
            f"({len(projection_channels_list)})."
        context = SPPM(in_channels=backbone_out_channels[-1],
                       inter_channels=sppm_inter_channels,
                       out_channels=sppm_out_channels,
                       pool_sizes=sppm_pool_sizes,
                       upsample_mode=sppm_upsample_mode,
                       align_corners=align_corners)
        self.encoder = PPLiteSegEncoder(backbone=backbone,
                                        context_module=context,
                                        projection_channels_list=projection_channels_list)
        encoder_channels = self.encoder.get_output_number_of_channels()

        # Init Decoder
        self.decoder = PPLiteSegDecoder(encoder_channels=encoder_channels,
                                        up_factors=decoder_up_factors,
                                        out_channels=decoder_channels,
                                        upsample_mode=decoder_upsample_mode,
                                        align_corners=align_corners)

        # Init Segmentation classification heads
        self.seg_head = nn.Sequential(
            SegmentationHead(in_channels=decoder_channels[-1],
                             mid_channels=head_mid_channels,
                             num_classes=num_classes,
                             dropout=dropout),
            make_upsample_module(scale_factor=head_scale_factor, upsample_mode=head_upsample_mode,
                                 align_corners=align_corners)
        )
        # Auxiliary heads
        if self.use_aux_heads:
            encoder_out_channels = projection_channels_list
            self.aux_heads = nn.ModuleList([
                nn.Sequential(
                    SegmentationHead(backbone_ch, hidden_ch, num_classes, dropout=dropout),
                    make_upsample_module(scale_factor=scale_factor, upsample_mode=head_upsample_mode,
                                         align_corners=align_corners)
                ) for backbone_ch, hidden_ch, scale_factor in zip(encoder_out_channels, aux_hidden_channels,
                                                                  aux_scale_factors)
            ])
        self.init_params()

    def _remove_auxiliary_heads(self):
        if hasattr(self, "aux_heads"):
            del self.aux_heads

    @property
    def backbone(self) -> nn.Module:
        """
        Support SG load backbone when training.
        """
        return self.encoder.backbone

    def forward(self, x):
        feats = self.encoder(x)
        if self.use_aux_heads:
            enc_feats = feats[:-1]
        x = self.decoder(feats)
        x = self.seg_head(x)
        if not self.use_aux_heads:
            return x
        aux_feats = [aux_head(feat) for feat, aux_head in zip(enc_feats, self.aux_heads)]
        return tuple([x] + aux_feats)

    def initialize_param_groups(self, lr: float, training_params: HpmStruct) -> list:
        """
        Custom param groups for training:
            - Different lr for backbone and the rest, if `multiply_head_lr` key is in `training_params`.
        """
        multiply_head_lr = get_param(training_params, "multiply_head_lr", 1)
        multiply_lr_params, no_multiply_params = self._separate_lr_multiply_params()
        param_groups = [{"named_params": no_multiply_params, "lr": lr, "name": "no_multiply_params"},
                        {"named_params": multiply_lr_params, "lr": lr * multiply_head_lr, "name": "multiply_lr_params"}]
        return param_groups

    def update_param_groups(self, param_groups: list, lr: float, epoch: int, iter: int, training_params: HpmStruct,
                            total_batch: int) -> list:
        multiply_head_lr = get_param(training_params, "multiply_head_lr", 1)
        for param_group in param_groups:
            param_group['lr'] = lr
            if param_group["name"] == "multiply_lr_params":
                param_group['lr'] *= multiply_head_lr
        return param_groups

    def _separate_lr_multiply_params(self):
        """
        Separate backbone params from the rest.
        :return: iterators of groups named_parameters.
        """
        multiply_lr_params, no_multiply_params = {}, {}
        for name, param in self.named_parameters():
            if "encoder.backbone" in name:
                no_multiply_params[name] = param
            else:
                multiply_lr_params[name] = param
        return multiply_lr_params.items(), no_multiply_params.items()

    def prep_model_for_conversion(self, input_size: Union[tuple, list], stride_ratio: int = 32, **kwargs):
        super().prep_model_for_conversion(input_size, **kwargs)
        if isinstance(self.encoder.context_module, SPPM):
            self.encoder.context_module.prep_model_for_conversion(input_size=input_size, stride_ratio=stride_ratio)


class PPLiteSegB(PPLiteSegBase):
    def __init__(self, arch_params: HpmStruct):
        backbone = STDC2Backbone(in_channels=get_param(arch_params, "in_channels", 3),
                                 out_down_ratios=[8, 16, 32])
        super().__init__(num_classes=get_param(arch_params, "num_classes"),
                         backbone=backbone,
                         projection_channels_list=[96, 128, 128],
                         sppm_inter_channels=128,
                         sppm_out_channels=128,
                         sppm_pool_sizes=[1, 2, 4],
                         sppm_upsample_mode="bilinear",
                         align_corners=False,
                         decoder_up_factors=[1, 2, 2],
                         decoder_channels=[128, 96, 64],
                         decoder_upsample_mode="bilinear",
                         head_scale_factor=8,
                         head_upsample_mode="bilinear",
                         head_mid_channels=64,
                         dropout=get_param(arch_params, "dropout", 0.),
                         use_aux_heads=get_param(arch_params, "use_aux_heads", False),
                         aux_hidden_channels=[32, 64, 64],
                         aux_scale_factors=[8, 16, 32])


class PPLiteSegT(PPLiteSegBase):
    def __init__(self, arch_params: HpmStruct):
        backbone = STDC1Backbone(in_channels=get_param(arch_params, "in_channels", 3),
                                 out_down_ratios=[8, 16, 32])
        super().__init__(num_classes=get_param(arch_params, "num_classes"),
                         backbone=backbone,
                         projection_channels_list=[64, 128, 128],
                         sppm_inter_channels=128,
                         sppm_out_channels=128,
                         sppm_pool_sizes=[1, 2, 4],
                         sppm_upsample_mode="bilinear",
                         align_corners=False,
                         decoder_up_factors=[1, 2, 2],
                         decoder_channels=[128, 64, 32],
                         decoder_upsample_mode="bilinear",
                         head_scale_factor=8,
                         head_upsample_mode="bilinear",
                         head_mid_channels=32,
                         dropout=get_param(arch_params, "dropout", 0.),
                         use_aux_heads=get_param(arch_params, "use_aux_heads", False),
                         aux_hidden_channels=[32, 64, 64],
                         aux_scale_factors=[8, 16, 32])

from super_gradients.training.models.sg_module import SgModule
import torch.nn as nn
from abc import abstractmethod, ABC


class SegmentationModule(SgModule, ABC):
    """
    Base SegmentationModule class
    """
    def __init__(self, use_aux_heads: bool):
        super().__init__()
        self._use_aux_heads = use_aux_heads

    @property
    def use_aux_heads(self):
        return self._use_aux_heads

    @use_aux_heads.setter
    def use_aux_heads(self, use_aux: bool):
        """
        public setter for self._use_aux_heads, called every time an assignment to self.use_aux_heads is applied.
        if use_aux is False, `_remove_auxiliary_heads` is called to delete auxiliary and detail heads.
        if use_aux is True, and self._use_aux_heads was already set to False a ValueError is raised, recreating
            aux and detail heads outside init method is not allowed, and the module should be recreated.
        """
        if use_aux is True and self._use_aux_heads is False:
            raise ValueError("Cant turn use_aux_heads from False to True. Try initiating the module again with"
                             " `use_aux_heads=True` or initiating the auxiliary heads modules manually.")
        if not use_aux:
            self._remove_auxiliary_heads()
        self._use_aux_heads = use_aux

    def prep_model_for_conversion(self, *args, **kwargs):
        # set to false and delete auxiliary and detail heads modules.
        self.use_aux_heads = False

    @abstractmethod
    def _remove_auxiliary_heads(self):
        raise NotImplementedError()

    @property
    @abstractmethod
    def backbone(self) -> nn.Module:
        """
        For SgModel load_backbone compatibility.
        """
        raise NotImplementedError()

    def init_params(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity="relu")
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

@@ -39,8 +39,12 @@ MODEL_URLS = {"regnetY800_imagenet": "https://deci-pretrained-models.s3.amazonaw
               "yolox_m_coco": "https://deci-pretrained-models.s3.amazonaws.com/yolox_coco/yolox_m_coco/average_model.pth",
               "yolox_l_coco": "https://deci-pretrained-models.s3.amazonaws.com/yolox_coco/yolox_l_coco/average_model.pth",
               "yolox_t_coco": "https://deci-pretrained-models.s3.amazonaws.com/yolox_coco/yolox_tiny_coco/ckpt_best.pth",
-              "yolox_n_coco": "https://deci-pretrained-models.s3.amazonaws.com/yolox_coco/yolox_n_coco/ckpt_best.pth"
+              "yolox_n_coco": "https://deci-pretrained-models.s3.amazonaws.com/yolox_coco/yolox_n_coco/ckpt_best.pth",
 
+              "pp_lite_t_seg50_cityscapes": "https://deci-pretrained-models.s3.amazonaws.com/ppliteseg/cityscapes/pplite_t_seg50/average_model.pth",
+              "pp_lite_t_seg75_cityscapes": "https://deci-pretrained-models.s3.amazonaws.com/ppliteseg/cityscapes/pplite_t_seg75/average_model.pth",
+              "pp_lite_b_seg50_cityscapes": "https://deci-pretrained-models.s3.amazonaws.com/ppliteseg/cityscapes/pplite_b_seg50/average_model.pth",
+              "pp_lite_b_seg75_cityscapes": "https://deci-pretrained-models.s3.amazonaws.com/ppliteseg/cityscapes/pplite_b_seg75/average_model.pth",
               }
 
 PRETRAINED_NUM_CLASSES = {"imagenet": 1000,

@@ -1,9 +1,11 @@
 from collections import OrderedDict
 import copy
-from typing import List, Union, Tuple
+from typing import List, Union, Tuple, Optional
 import torch
 from torch import nn
 
+from super_gradients.common import UpsampleMode
+
 
 class MultiOutputModule(nn.Module):
     """
@@ -219,3 +221,23 @@ class NormalizationAdapter(torch.nn.Module):
     def forward(self, x):
         x = (x + self.additive) * self.multiplier
         return x
+
+
+def make_upsample_module(scale_factor: int,
+                         upsample_mode: Union[str, UpsampleMode],
+                         align_corners: Optional[bool] = None):
+    """
+    Factory method for creating upsampling modules.
+    :param scale_factor: upsample scale factor
+    :param upsample_mode: see UpsampleMode for supported options.
+    :return: nn.Module
+    """
+    upsample_mode = upsample_mode.value if isinstance(upsample_mode, UpsampleMode) else upsample_mode
+    if upsample_mode == UpsampleMode.NEAREST.value:
+        # Prevent ValueError when passing align_corners with nearest mode.
+        module = nn.Upsample(scale_factor=scale_factor, mode=upsample_mode)
+    elif upsample_mode in [UpsampleMode.BILINEAR.value, UpsampleMode.BICUBIC.value]:
+        module = nn.Upsample(scale_factor=scale_factor, mode=upsample_mode, align_corners=align_corners)
+    else:
+        raise NotImplementedError(f"Upsample mode: `{upsample_mode}` is not supported.")
+    return module

@@ -11,7 +11,6 @@ import os
 import shutil
 from super_gradients.training.utils.ssd_utils import SSDPostPredictCallback
 from super_gradients.training.models.detection_models.ssd import DEFAULT_SSD_LITE_MOBILENET_V2_ARCH_PARAMS
-import torchvision.transforms as transforms
 from super_gradients.training.losses.ddrnet_loss import DDRNetLoss
 from super_gradients.training.metrics import DetectionMetrics
 from super_gradients.training.transforms.transforms import Rescale
@@ -232,7 +231,9 @@ class PretrainedModelsTest(unittest.TestCase):
         self.cityscapes_pretrained_arch_params = {
             "ddrnet_23": {"aux_head": True, "sync_bn": True},
             "regseg48": {},
-            "stdc": {"use_aux_heads": True, "aux_head": True}}
+            "stdc": {"use_aux_heads": True, "aux_head": True},
+            "pplite_seg": {"use_aux_heads": True},
+        }
 
         self.cityscapes_pretrained_ckpt_params = {"pretrained_weights": "cityscapes"}
         self.cityscapes_pretrained_mious = {"ddrnet_23": 0.8026,
@@ -241,7 +242,11 @@ class PretrainedModelsTest(unittest.TestCase):
                                             "stdc1_seg75": 0.7687,
                                             "stdc2_seg50": 0.7644,
                                             "stdc2_seg75": 0.7893,
-                                            "regseg48": 0.7815}
+                                            "regseg48": 0.7815,
+                                            "pp_lite_t_seg50": 0.7492,
+                                            "pp_lite_t_seg75": 0.7756,
+                                            "pp_lite_b_seg50": 0.7648,
+                                            "pp_lite_b_seg75": 0.7852}
 
         self.cityscapes_dataset = CityscapesDatasetInterface(dataset_params={
             "batch_size": 3,
@@ -249,22 +254,22 @@ class PretrainedModelsTest(unittest.TestCase):
             "dataset_dir": "/data/cityscapes/",
             "crop_size": 1024,
             "img_size": 1024,
-            "image_mask_transforms_aug": transforms.Compose([]),
-            "image_mask_transforms": transforms.Compose([])  # no transform for evaluation
+            "image_mask_transforms_aug": [],
+            "image_mask_transforms": []  # no transform for evaluation
         }, cache_labels=False)
 
         self.cityscapes_dataset_rescaled50 = CityscapesDatasetInterface(dataset_params={
             "batch_size": 3,
             "val_batch_size": 3,
-            "image_mask_transforms_aug": transforms.Compose([]),
-            "image_mask_transforms": transforms.Compose([Rescale(scale_factor=0.5)])  # no transform for evaluation
+            "image_mask_transforms_aug": [],
+            "image_mask_transforms": [Rescale(scale_factor=0.5)]  # no transform for evaluation
         }, cache_labels=False)
 
         self.cityscapes_dataset_rescaled75 = CityscapesDatasetInterface(dataset_params={
             "batch_size": 3,
             "val_batch_size": 3,
-            "image_mask_transforms_aug": transforms.Compose([]),
-            "image_mask_transforms": transforms.Compose([Rescale(scale_factor=0.75)])  # no transform for evaluation
+            "image_mask_transforms_aug": [],
+            "image_mask_transforms": [Rescale(scale_factor=0.75)]  # no transform for evaluation
         }, cache_labels=False)
 
         self.transfer_segmentation_dataset = SegmentationTestDatasetInterface(image_size=1024)
@@ -788,6 +793,50 @@ class PretrainedModelsTest(unittest.TestCase):
                            metrics_progress_verbose=True)[0].cpu().item()
         self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc2_seg75"], delta=0.001)
 
+    def test_pretrained_pplite_t_seg50_cityscapes(self):
+        trainer = Trainer('cityscapes_pretrained_pplite_t_seg50', model_checkpoints_location='local',
+                          multi_gpu=MultiGPUMode.OFF)
+        trainer.connect_dataset_interface(self.cityscapes_dataset_rescaled50, data_loader_num_workers=8)
+        trainer.build_model("pp_lite_t_seg50", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        res = trainer.test(test_loader=self.cityscapes_dataset_rescaled50.val_loader,
+                           test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
+                           metrics_progress_verbose=True)[0].cpu().item()
+        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_t_seg50"], delta=0.001)
+
+    def test_pretrained_pplite_t_seg75_cityscapes(self):
+        trainer = Trainer('cityscapes_pretrained_pplite_t_seg75', model_checkpoints_location='local',
+                          multi_gpu=MultiGPUMode.OFF)
+        trainer.connect_dataset_interface(self.cityscapes_dataset_rescaled75, data_loader_num_workers=8)
+        trainer.build_model("pp_lite_t_seg75", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        res = trainer.test(test_loader=self.cityscapes_dataset_rescaled75.val_loader,
+                           test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
+                           metrics_progress_verbose=True)[0].cpu().item()
+        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_t_seg75"], delta=0.001)
+
+    def test_pretrained_pplite_b_seg50_cityscapes(self):
+        trainer = Trainer('cityscapes_pretrained_pplite_b_seg50', model_checkpoints_location='local',
+                          multi_gpu=MultiGPUMode.OFF)
+        trainer.connect_dataset_interface(self.cityscapes_dataset_rescaled50, data_loader_num_workers=8)
+        trainer.build_model("pp_lite_b_seg50", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        res = trainer.test(test_loader=self.cityscapes_dataset_rescaled50.val_loader,
+                           test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
+                           metrics_progress_verbose=True)[0].cpu().item()
+        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_b_seg50"], delta=0.001)
+
+    def test_pretrained_pplite_b_seg75_cityscapes(self):
+        trainer = Trainer('cityscapes_pretrained_pplite_b_seg75', model_checkpoints_location='local',
+                          multi_gpu=MultiGPUMode.OFF)
+        trainer.connect_dataset_interface(self.cityscapes_dataset_rescaled75, data_loader_num_workers=8)
+        trainer.build_model("pp_lite_b_seg75", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"],
+                            checkpoint_params=self.cityscapes_pretrained_ckpt_params)
+        res = trainer.test(test_loader=self.cityscapes_dataset_rescaled75.val_loader,
+                           test_metrics_list=[IoU(num_classes=20, ignore_index=19)],
+                           metrics_progress_verbose=True)[0].cpu().item()
+        self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_b_seg75"], delta=0.001)
+
     def test_transfer_learning_stdc2_seg75_cityscapes(self):
         trainer = Trainer('cityscapes_pretrained_stdc2_seg75_transfer_learning', model_checkpoints_location='local',
                           multi_gpu=MultiGPUMode.OFF)