Source code for torchvision.models.segmentation.lraspp

from collections import OrderedDict
from functools import partial
from typing import Any, Dict, Optional

from torch import nn, Tensor
from torch.nn import functional as F

from ...transforms._presets import SemanticSegmentation
from ...utils import _log_api_usage_once
from .._api import register_model, Weights, WeightsEnum
from .._meta import _VOC_CATEGORIES
from .._utils import _ovewrite_value_param, handle_legacy_interface, IntermediateLayerGetter
from ..mobilenetv3 import mobilenet_v3_large, MobileNet_V3_Large_Weights, MobileNetV3

__all__ = ["LRASPP", "LRASPP_MobileNet_V3_Large_Weights", "lraspp_mobilenet_v3_large"]

class LRASPP(nn.Module):
    Implements a Lite R-ASPP Network for semantic segmentation from
    `"Searching for MobileNetV3"

        backbone (nn.Module): the network used to compute the features for the model.
            The backbone should return an OrderedDict[Tensor], with the key being
            "high" for the high level feature map and "low" for the low level feature map.
        low_channels (int): the number of channels of the low level features.
        high_channels (int): the number of channels of the high level features.
        num_classes (int, optional): number of output classes of the model (including the background).
        inter_channels (int, optional): the number of channels for intermediate computations.

    def __init__(
        self, backbone: nn.Module, low_channels: int, high_channels: int, num_classes: int, inter_channels: int = 128
    ) -> None:
        self.backbone = backbone
        self.classifier = LRASPPHead(low_channels, high_channels, num_classes, inter_channels)

    def forward(self, input: Tensor) -> Dict[str, Tensor]:
        features = self.backbone(input)
        out = self.classifier(features)
        out = F.interpolate(out, size=input.shape[-2:], mode="bilinear", align_corners=False)

        result = OrderedDict()
        result["out"] = out

        return result

class LRASPPHead(nn.Module):
    def __init__(self, low_channels: int, high_channels: int, num_classes: int, inter_channels: int) -> None:
        self.cbr = nn.Sequential(
            nn.Conv2d(high_channels, inter_channels, 1, bias=False),
        self.scale = nn.Sequential(
            nn.Conv2d(high_channels, inter_channels, 1, bias=False),
        self.low_classifier = nn.Conv2d(low_channels, num_classes, 1)
        self.high_classifier = nn.Conv2d(inter_channels, num_classes, 1)

    def forward(self, input: Dict[str, Tensor]) -> Tensor:
        low = input["low"]
        high = input["high"]

        x = self.cbr(high)
        s = self.scale(high)
        x = x * s
        x = F.interpolate(x, size=low.shape[-2:], mode="bilinear", align_corners=False)

        return self.low_classifier(low) + self.high_classifier(x)

def _lraspp_mobilenetv3(backbone: MobileNetV3, num_classes: int) -> LRASPP:
    backbone = backbone.features
    # Gather the indices of blocks which are strided. These are the locations of C1, ..., Cn-1 blocks.
    # The first and last blocks are always included because they are the C0 (conv1) and Cn.
    stage_indices = [0] + [i for i, b in enumerate(backbone) if getattr(b, "_is_cn", False)] + [len(backbone) - 1]
    low_pos = stage_indices[-4]  # use C2 here which has output_stride = 8
    high_pos = stage_indices[-1]  # use C5 which has output_stride = 16
    low_channels = backbone[low_pos].out_channels
    high_channels = backbone[high_pos].out_channels
    backbone = IntermediateLayerGetter(backbone, return_layers={str(low_pos): "low", str(high_pos): "high"})

    return LRASPP(backbone, low_channels, high_channels, num_classes)

[docs]class LRASPP_MobileNet_V3_Large_Weights(WeightsEnum): COCO_WITH_VOC_LABELS_V1 = Weights( url="", transforms=partial(SemanticSegmentation, resize_size=520), meta={ "num_params": 3221538, "categories": _VOC_CATEGORIES, "min_size": (1, 1), "recipe": "", "_metrics": { "COCO-val2017-VOC-labels": { "miou": 57.9, "pixel_acc": 91.2, } }, "_ops": 2.086, "_file_size": 12.49, "_docs": """ These weights were trained on a subset of COCO, using only the 20 categories that are present in the Pascal VOC dataset. """, }, ) DEFAULT = COCO_WITH_VOC_LABELS_V1
[docs]@register_model() @handle_legacy_interface( weights=("pretrained", LRASPP_MobileNet_V3_Large_Weights.COCO_WITH_VOC_LABELS_V1), weights_backbone=("pretrained_backbone", MobileNet_V3_Large_Weights.IMAGENET1K_V1), ) def lraspp_mobilenet_v3_large( *, weights: Optional[LRASPP_MobileNet_V3_Large_Weights] = None, progress: bool = True, num_classes: Optional[int] = None, weights_backbone: Optional[MobileNet_V3_Large_Weights] = MobileNet_V3_Large_Weights.IMAGENET1K_V1, **kwargs: Any, ) -> LRASPP: """Constructs a Lite R-ASPP Network model with a MobileNetV3-Large backbone from `Searching for MobileNetV3 <>`_ paper. .. betastatus:: segmentation module Args: weights (:class:`~torchvision.models.segmentation.LRASPP_MobileNet_V3_Large_Weights`, optional): The pretrained weights to use. See :class:`~torchvision.models.segmentation.LRASPP_MobileNet_V3_Large_Weights` below for more details, and possible values. By default, no pre-trained weights are used. progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. num_classes (int, optional): number of output classes of the model (including the background). aux_loss (bool, optional): If True, it uses an auxiliary loss. weights_backbone (:class:`~torchvision.models.MobileNet_V3_Large_Weights`, optional): The pretrained weights for the backbone. **kwargs: parameters passed to the ``torchvision.models.segmentation.LRASPP`` base class. Please refer to the `source code <>`_ for more details about this class. .. autoclass:: torchvision.models.segmentation.LRASPP_MobileNet_V3_Large_Weights :members: """ if kwargs.pop("aux_loss", False): raise NotImplementedError("This model does not use auxiliary loss") weights = LRASPP_MobileNet_V3_Large_Weights.verify(weights) weights_backbone = MobileNet_V3_Large_Weights.verify(weights_backbone) if weights is not None: weights_backbone = None num_classes = _ovewrite_value_param("num_classes", num_classes, len(weights.meta["categories"])) elif num_classes is None: num_classes = 21 backbone = mobilenet_v3_large(weights=weights_backbone, dilated=True) model = _lraspp_mobilenetv3(backbone, num_classes) if weights is not None: model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True)) return model


Access comprehensive developer documentation for PyTorch

View Docs


Get in-depth tutorials for beginners and advanced developers

View Tutorials


Find development resources and get your questions answered

View Resources