fcn_101代码理解

def _segm_resnet(name, backbone_name, num_classes, aux, pretrained_backbone=True):
"""
name = fcn,
backbone_name=resnet101, 
num_classes=21
"""
# 基本的resnet101
    backbone = resnet.__dict__[backbone_name](
        pretrained=pretrained_backbone,
        replace_stride_with_dilation=[False, True, True])
	# 要返回的层
    return_layers = {'layer4': 'out'}
    if aux:  # Tue，则return_layers字典添加
        return_layers['layer3'] = 'aux'
    backbone = IntermediateLayerGetter(backbone, return_layers=return_layers)

    aux_classifier = None
    if aux:  # 层减半的分类层
        inplanes = 1024
        aux_classifier = FCNHead(inplanes, num_classes)

    model_map = {
        'deeplabv3': (DeepLabHead, DeepLabV3),
        'fcn': (FCNHead, FCN),
    }
    inplanes = 2048
    #  model_map[‘fcn’][FCNHead] 函数，见下面FCNHead类
    # 分类层 21类
    classifier = model_map[name][0](inplanes, num_classes)
    # 特征提取层
    # model_map[‘fcn’][FCN] 函数
    # <class 'torchvision.models.segmentation.fcn.FCN'>
    base_model = model_map[name][1]  
   # 构建模型
    model = base_model(backbone, classifier, aux_classifier)
    return model

class FCNHead(nn.Sequential):
    def __init__(self, in_channels, channels):
        inter_channels = in_channels // 4
        layers = [
            nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
            nn.BatchNorm2d(inter_channels),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Conv2d(inter_channels, channels, 1)
        ]

        super(FCNHead, self).__init__(*layers)

class FCN(_SimpleSegmentationModel):
# 继承_SimpleSegmentationModel类， 见下_SimpleSegmentationModel类
    """
    Implements a Fully-Convolutional Network for semantic segmentation.

    Arguments:
        backbone (nn.Module): the network used to compute the features for the model.
            The backbone should return an OrderedDict[Tensor], with the key being
            "out" for the last feature map used, and "aux" if an auxiliary classifier
            is used.
        classifier (nn.Module): module that takes the "out" element returned from
            the backbone and returns a dense prediction.
        aux_classifier (nn.Module, optional): auxiliary classifier used during training
    """
    pass

class _SimpleSegmentationModel(nn.Module):
# 最基础的模型定义
    __constants__ = ['aux_classifier']

    def __init__(self, backbone, classifier, aux_classifier=None):
        super(_SimpleSegmentationModel, self).__init__()
        self.backbone = backbone
        self.classifier = classifier
        self.aux_classifier = aux_classifier

    def forward(self, x):
        input_shape = x.shape[-2:]  # torch.Size([224, 224])
        # contract: features is a dict of tensors
        features = self.backbone(x)
        # OrderedDict([('out', torch.Size([1, 2048, 28, 28])),'aux',torch.Size([1, 1024, 28, 28])])

        result = OrderedDict()  # 定义result字典
        x = features["out"]  # torch.Size([1, 2048, 28, 28])
        x = self.classifier(x)  # torch.Size([1, 21, 28, 28])
        x = F.interpolate(x, size=input_shape, mode='bilinear', align_corners=False) # 上采样到原图大小torch.Size([1, 21, 224, 224])
        result["out"] = x  # 赋值输出结果

        if self.aux_classifier is not None: #如果aux存在，同上操作，得到结果
            x = features["aux"]  # torch.Size([1, 1024, 28, 28])
            x = self.aux_classifier(x) torch.Size([1, 21, 28, 28])
            x = F.interpolate(x, size=input_shape, mode='bilinear', align_corners=False)
            result["aux"] = x

        return result

class IntermediateLayerGetter(nn.ModuleDict):
    """
    Module wrapper that returns intermediate layers from a model

    It has a strong assumption that the modules have been registered
    into the model in the same order as they are used.
    This means that one should **not** reuse the same nn.Module
    twice in the forward if you want this to work.

    Additionally, it is only able to query submodules that are directly
    assigned to the model. So if `model` is passed, `model.feature1` can
    be returned, but not `model.feature1.layer2`.

    Arguments:
        model (nn.Module): model on which we will extract the features
        return_layers (Dict[name, new_name]): a dict containing the names
            of the modules for which the activations will be returned as
            the key of the dict, and the value of the dict is the name
            of the returned activation (which the user can specify).

    Examples::

        >>> m = torchvision.models.resnet18(pretrained=True)
        >>> # extract layer1 and layer3, giving as names `feat1` and feat2`
        >>> new_m = torchvision.models._utils.IntermediateLayerGetter(m,
        >>>     {'layer1': 'feat1', 'layer3': 'feat2'})
        >>> out = new_m(torch.rand(1, 3, 224, 224))
        >>> print([(k, v.shape) for k, v in out.items()])
        >>>     [('feat1', torch.Size([1, 64, 56, 56])),
        >>>      ('feat2', torch.Size([1, 256, 14, 14]))]
    """
    _version = 2
    __annotations__ = {
        "return_layers": Dict[str, str],
    }

    def __init__(self, model, return_layers):
        if not set(return_layers).issubset([name for name, _ in model.named_children()]):
            raise ValueError("return_layers are not present in model")
        orig_return_layers = return_layers
        return_layers = {str(k): str(v) for k, v in return_layers.items()}
        layers = OrderedDict()
        for name, module in model.named_children():
            layers[name] = module
            if name in return_layers:
                del return_layers[name]
            if not return_layers:
                break

        super(IntermediateLayerGetter, self).__init__(layers)
        self.return_layers = orig_return_layers
        # <class 'dict'>: {'layer4': 'out'} 返回layer4层的结果， 将输出的字典key取值为out

    def forward(self, x):
        out = OrderedDict()
        for name, module in self.items():  # 遍历字典中的每一层函数
        # name= 'conv1'，module =Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
            x = module(x)  # 执行
            if name in self.return_layers: # 如果key在返回层中
            # 'layer4' in 
                out_name = self.return_layers[name]  # 返回该层对应的输出名字'out'
                out[out_name] = x  赋值
        return out