Skip to content

Latest commit

 

History

History
257 lines (202 loc) · 6.98 KB

Add_modules.md

File metadata and controls

257 lines (202 loc) · 6.98 KB

增加新的模块

BILIBILI

  • 一个完整的模型由Backbone、Neck、Head、Loss组成，在文件夹configs下可以找到
  • 主干网络：通常是一个特征提取网络，例如 ResNet、MobileNet
  • 颈部：用于连接主干网络和头部的组件，例如 GlobalAveragePooling
  • 头部：用于执行特定任务的组件，例如分类和回归
  • 损失：用于计算预测值与真实值偏差值

添加新的主干网络Backbone

以 `ResNet_CIFAR` 为例

ResNet_CIFAR 针对 CIFAR 32x32 的图像输入，将 ResNet 中 kernel_size=7, stride=2 的设置替换为 kernel_size=3, stride=1，并移除了 stem 层之后的 MaxPooling，以避免传递过小的特征图到残差块中。

它继承自 ResNet 并只修改了 stem 层。
  1. 创建一个新文件 configs/backbones/resnet_cifar.py。
import torch.nn as nn

from ..common import BaseModule
from .resnet import ResNet


class ResNet_CIFAR(ResNet):

    """ResNet backbone for CIFAR.

    （对这个主干网络的简短描述）

    Args:
        depth(int): Network depth, from {18, 34, 50, 101, 152}.
        ...
        （参数文档）
    """

    def __init__(self, depth, deep_stem=False, **kwargs):
        # 调用基类 ResNet 的初始化函数
        super(ResNet_CIFAR, self).__init__(depth, deep_stem=deep_stem **kwargs)
        # 其他特殊的初始化流程
        assert not self.deep_stem, 'ResNet_CIFAR do not support deep_stem'

    def _make_stem_layer(self, in_channels, base_channels):
        # 重载基类的方法，以实现对网络结构的修改
        self.conv1 = build_conv_layer(
            self.conv_cfg,
            in_channels,
            base_channels,
            kernel_size=3,
            stride=1,
            padding=1,
            bias=False)
        self.norm1_name, norm1 = build_norm_layer(
            self.norm_cfg, base_channels, postfix=1)
        self.add_module(self.norm1_name, norm1)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):  # 需要返回一个元组
        pass  # 此处省略了网络的前向实现

    def init_weights(self, pretrained=None):
        pass  # 如果有必要的话，重载基类 ResNet 的参数初始化函数

    def train(self, mode=True):
        pass  # 如果有必要的话，重载基类 ResNet 的训练状态函数
  1. 在 configs/backbones/__init__.py 中导入新模块
...
from .resnet_cifar import ResNet_CIFAR

__all__ = [
    ..., 'ResNet_CIFAR'
]
  1. 在配置文件中使用新的主干网络
model_cfg = dict(
    backbone=dict(
        type='ResNet_CIFAR',
        depth=18,
        other_arg=xxx),
    ...

添加新的颈部Neck

以 `GlobalAveragePooling` 为例

要添加新的颈部组件，主要需要实现 forward 函数，该函数对主干网络的输出进行 一些操作并将结果传递到头部。
  1. 创建一个新文件configs/necks/gap.py
import torch.nn as nn


class GlobalAveragePooling(nn.Module):

    def __init__(self):
        self.gap = nn.AdaptiveAvgPool2d((1, 1))

    def forward(self, inputs):
        # 简单起见，我们默认输入是一个张量
        outs = self.gap(inputs)
        outs = outs.view(inputs.size(0), -1)
        return outs
  1. 在configs/necks/__init__.py 中导入新模块
...
from .gap import GlobalAveragePooling

__all__ = [
    ..., 'GlobalAveragePooling'
]
  1. 修改配置文件以使用新的颈部组件
model_cfg = dict(
    neck=dict(type='GlobalAveragePooling'),
)

添加新的头部Head

以 `LinearClsHead` 为例

要添加新的颈部组件，主要需要实现 forward 函数，该函数对主干网络的输出进行 一些操作并将结果传递到头部。
  1. 创建一个新文件configs/heads/linear_head.py
from .cls_head import ClsHead

class LinearClsHead(ClsHead):

    def __init__(self,
              num_classes,
              in_channels,
              loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
              topk=(1, )):
        super(LinearClsHead, self).__init__(loss=loss, topk=topk)
        self.in_channels = in_channels
        self.num_classes = num_classes

        if self.num_classes <= 0:
            raise ValueError(
                f'num_classes={num_classes} must be a positive integer')

        self._init_layers()

    def _init_layers(self):
        self.fc = nn.Linear(self.in_channels, self.num_classes)

    def init_weights(self):
        normal_init(self.fc, mean=0, std=0.01, bias=0)

    def forward_train(self, x, gt_label):
        cls_score = self.fc(x)
        losses = self.loss(cls_score, gt_label)
        return losses
  1. 在configs/heads/__init__.py 中导入新模块
...
from .linear_head import LinearClsHead

__all__ = [
    ..., 'LinearClsHead'
]
  1. 修改配置文件以使用新的颈部组件，连同 GlobalAveragePooling 颈部组件，完整的模型配置如下：
model_cfg = dict(
    backbone=dict(
        type='ResNet',
        depth=50,
        num_stages=4,
        out_indices=(3, ),
        style='pytorch'),
    neck=dict(type='GlobalAveragePooling'),
    head=dict(
        type='LinearClsHead',
        num_classes=1000,
        in_channels=2048,
        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
        topk=(1, 5),
    ))

添加新的损失函数Loss

要添加新的损失函数，主要需要在损失函数模块中 forward 函数。另外，利用装饰器 weighted_loss 可以方便的实现对每个元素的损失进行加权平均。

假设我们要模拟从另一个分类模型生成的概率分布，需要添加 L1loss 来实现该目的。
  1. 创建一个新文件configs/losses/l1_loss.py
import torch
import torch.nn as nn

from .utils import weighted_loss

@weighted_loss
def l1_loss(pred, target):
    assert pred.size() == target.size() and target.numel() > 0
    loss = torch.abs(pred - target)
    return loss

class L1Loss(nn.Module):

    def __init__(self, reduction='mean', loss_weight=1.0):
        super(L1Loss, self).__init__()
        self.reduction = reduction
        self.loss_weight = loss_weight

    def forward(self,
                pred,
                target,
                weight=None,
                avg_factor=None,
                reduction_override=None):
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = (
            reduction_override if reduction_override else self.reduction)
        loss = self.loss_weight * l1_loss(
            pred, target, weight, reduction=reduction, avg_factor=avg_factor)
        return loss
  1. 在configs/losses/__init__.py 中导入新模块
...
from .l1_loss import L1Loss, l1_loss

__all__ = [
    ..., 'L1Loss', 'l1_loss'
]
  1. 修改配置文件中的 loss 字段以使用新的损失函数
loss=dict(type='L1Loss', loss_weight=1.0))