|
|
|
|
|
from typing import Optional, Sequence
|
|
|
|
|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
import torch.utils.checkpoint as cp
|
|
|
from mmcv.cnn import ConvModule
|
|
|
from mmengine.logging import MMLogger
|
|
|
from mmengine.model.weight_init import constant_init, kaiming_init
|
|
|
from mmengine.runner import load_checkpoint
|
|
|
from mmengine.utils.dl_utils.parrots_wrapper import _BatchNorm
|
|
|
from torch.nn.modules.utils import _ntuple
|
|
|
|
|
|
from mmaction.registry import MODELS
|
|
|
from mmaction.utils import ConfigType
|
|
|
|
|
|
|
|
|
class Bottleneck2dAudio(nn.Module):
|
|
|
"""Bottleneck2D block for ResNet2D.
|
|
|
|
|
|
Args:
|
|
|
inplanes (int): Number of channels for the input in first conv3d layer.
|
|
|
planes (int): Number of channels produced by some norm/conv3d layers.
|
|
|
stride (int): Stride in the conv layer. Defaults to 2.
|
|
|
dilation (int): Spacing between kernel elements. Defaults to 1.
|
|
|
downsample (nn.Module, optional): Downsample layer. Defaults to None.
|
|
|
factorize (bool): Whether to factorize kernel. Defaults to True.
|
|
|
norm_cfg (dict): Config for norm layers. required keys are ``type`` and
|
|
|
``requires_grad``. Defaults to None.
|
|
|
with_cp (bool): Use checkpoint or not. Using checkpoint will save some
|
|
|
memory while slowing down the trgaining speed. Defaults to False.
|
|
|
"""
|
|
|
expansion = 4
|
|
|
|
|
|
def __init__(self,
|
|
|
inplanes: int,
|
|
|
planes: int,
|
|
|
stride: int = 2,
|
|
|
dilation: int = 1,
|
|
|
downsample: Optional[nn.Module] = None,
|
|
|
factorize: bool = True,
|
|
|
norm_cfg: ConfigType = None,
|
|
|
with_cp: bool = False) -> None:
|
|
|
super().__init__()
|
|
|
|
|
|
self.inplanes = inplanes
|
|
|
self.planes = planes
|
|
|
self.stride = stride
|
|
|
self.dilation = dilation
|
|
|
self.factorize = factorize
|
|
|
self.norm_cfg = norm_cfg
|
|
|
self.with_cp = with_cp
|
|
|
|
|
|
self.conv1_stride = 1
|
|
|
self.conv2_stride = stride
|
|
|
|
|
|
conv1_kernel_size = (1, 1)
|
|
|
conv1_padding = 0
|
|
|
conv2_kernel_size = (3, 3)
|
|
|
conv2_padding = (dilation, dilation)
|
|
|
self.conv1 = ConvModule(
|
|
|
inplanes,
|
|
|
planes,
|
|
|
kernel_size=conv1_kernel_size,
|
|
|
padding=conv1_padding,
|
|
|
dilation=dilation,
|
|
|
norm_cfg=self.norm_cfg,
|
|
|
bias=False)
|
|
|
self.conv2 = ConvModule(
|
|
|
planes,
|
|
|
planes,
|
|
|
kernel_size=conv2_kernel_size,
|
|
|
stride=stride,
|
|
|
padding=conv2_padding,
|
|
|
dilation=dilation,
|
|
|
bias=False,
|
|
|
conv_cfg=dict(type='ConvAudio') if factorize else dict(
|
|
|
type='Conv'),
|
|
|
norm_cfg=None,
|
|
|
act_cfg=None)
|
|
|
self.conv3 = ConvModule(
|
|
|
2 * planes if factorize else planes,
|
|
|
planes * self.expansion,
|
|
|
kernel_size=1,
|
|
|
bias=False,
|
|
|
norm_cfg=self.norm_cfg,
|
|
|
act_cfg=None)
|
|
|
|
|
|
self.relu = nn.ReLU(inplace=True)
|
|
|
self.downsample = downsample
|
|
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
|
"""Defines the computation performed at every call.
|
|
|
|
|
|
Args:
|
|
|
x (torch.Tensor): The input data.
|
|
|
|
|
|
Returns:
|
|
|
torch.Tensor: The output of the module.
|
|
|
"""
|
|
|
|
|
|
def _inner_forward(x):
|
|
|
identity = x
|
|
|
out = self.conv1(x)
|
|
|
out = self.conv2(out)
|
|
|
out = self.conv3(out)
|
|
|
|
|
|
if self.downsample is not None:
|
|
|
identity = self.downsample(x)
|
|
|
out += identity
|
|
|
|
|
|
return out
|
|
|
|
|
|
if self.with_cp and x.requires_grad:
|
|
|
out = cp.checkpoint(_inner_forward, x)
|
|
|
else:
|
|
|
out = _inner_forward(x)
|
|
|
|
|
|
out = self.relu(out)
|
|
|
|
|
|
return out
|
|
|
|
|
|
|
|
|
@MODELS.register_module()
|
|
|
class ResNetAudio(nn.Module):
|
|
|
"""ResNet 2d audio backbone. Reference:
|
|
|
|
|
|
<https://arxiv.org/abs/2001.08740>`_.
|
|
|
|
|
|
Args:
|
|
|
depth (int): Depth of resnet, from ``{50, 101, 152}``.
|
|
|
pretrained (str, optional): Name of pretrained model. Defaults to None.
|
|
|
in_channels (int): Channel num of input features. Defaults to 1.
|
|
|
base_channels (int): Channel num of stem output features.
|
|
|
Defaults to 32.
|
|
|
num_stages (int): Resnet stages. Defaults to 4.
|
|
|
strides (Sequence[int]): Strides of residual blocks of each stage.
|
|
|
Defaults to ``(1, 2, 2, 2)``.
|
|
|
dilations (Sequence[int]): Dilation of each stage.
|
|
|
Defaults to ``(1, 1, 1, 1)``.
|
|
|
conv1_kernel (int): Kernel size of the first conv layer. Defaults to 9.
|
|
|
conv1_stride (Union[int, Tuple[int]]): Stride of the first conv layer.
|
|
|
Defaults to 1.
|
|
|
frozen_stages (int): Stages to be frozen (all param fixed). -1 means
|
|
|
not freezing any parameters. Defaults to -1.
|
|
|
factorize (Sequence[int]): factorize Dims of each block for audio.
|
|
|
Defaults to ``(1, 1, 0, 0)``.
|
|
|
norm_eval (bool): Whether to set BN layers to eval mode, namely, freeze
|
|
|
running stats (mean and var). Defaults to False.
|
|
|
with_cp (bool): Use checkpoint or not. Using checkpoint will save some
|
|
|
memory while slowing down the training speed. Defaults to False.
|
|
|
conv_cfg (Union[dict, ConfigDict]): Config for norm layers.
|
|
|
Defaults to ``dict(type='Conv')``.
|
|
|
norm_cfg (Union[dict, ConfigDict]): Config for norm layers. required
|
|
|
keys are ``type`` and ``requires_grad``.
|
|
|
Defaults to ``dict(type='BN2d', requires_grad=True)``.
|
|
|
act_cfg (Union[dict, ConfigDict]): Config for activate layers.
|
|
|
Defaults to ``dict(type='ReLU', inplace=True)``.
|
|
|
zero_init_residual (bool): Whether to use zero initialization
|
|
|
for residual block. Defaults to True.
|
|
|
"""
|
|
|
|
|
|
arch_settings = {
|
|
|
|
|
|
|
|
|
50: (Bottleneck2dAudio, (3, 4, 6, 3)),
|
|
|
101: (Bottleneck2dAudio, (3, 4, 23, 3)),
|
|
|
152: (Bottleneck2dAudio, (3, 8, 36, 3))
|
|
|
}
|
|
|
|
|
|
def __init__(self,
|
|
|
depth: int,
|
|
|
pretrained: str = None,
|
|
|
in_channels: int = 1,
|
|
|
num_stages: int = 4,
|
|
|
base_channels: int = 32,
|
|
|
strides: Sequence[int] = (1, 2, 2, 2),
|
|
|
dilations: Sequence[int] = (1, 1, 1, 1),
|
|
|
conv1_kernel: int = 9,
|
|
|
conv1_stride: int = 1,
|
|
|
frozen_stages: int = -1,
|
|
|
factorize: Sequence[int] = (1, 1, 0, 0),
|
|
|
norm_eval: bool = False,
|
|
|
with_cp: bool = False,
|
|
|
conv_cfg: ConfigType = dict(type='Conv'),
|
|
|
norm_cfg: ConfigType = dict(type='BN2d', requires_grad=True),
|
|
|
act_cfg: ConfigType = dict(type='ReLU', inplace=True),
|
|
|
zero_init_residual: bool = True) -> None:
|
|
|
super().__init__()
|
|
|
if depth not in self.arch_settings:
|
|
|
raise KeyError(f'invalid depth {depth} for resnet')
|
|
|
self.depth = depth
|
|
|
self.pretrained = pretrained
|
|
|
self.in_channels = in_channels
|
|
|
self.base_channels = base_channels
|
|
|
self.num_stages = num_stages
|
|
|
assert 1 <= num_stages <= 4
|
|
|
self.dilations = dilations
|
|
|
self.conv1_kernel = conv1_kernel
|
|
|
self.conv1_stride = conv1_stride
|
|
|
self.frozen_stages = frozen_stages
|
|
|
self.stage_factorization = _ntuple(num_stages)(factorize)
|
|
|
self.norm_eval = norm_eval
|
|
|
self.with_cp = with_cp
|
|
|
self.conv_cfg = conv_cfg
|
|
|
self.norm_cfg = norm_cfg
|
|
|
self.act_cfg = act_cfg
|
|
|
self.zero_init_residual = zero_init_residual
|
|
|
|
|
|
self.block, stage_blocks = self.arch_settings[depth]
|
|
|
self.stage_blocks = stage_blocks[:num_stages]
|
|
|
self.inplanes = self.base_channels
|
|
|
|
|
|
self._make_stem_layer()
|
|
|
|
|
|
self.res_layers = []
|
|
|
for i, num_blocks in enumerate(self.stage_blocks):
|
|
|
stride = strides[i]
|
|
|
dilation = dilations[i]
|
|
|
planes = self.base_channels * 2**i
|
|
|
res_layer = self.make_res_layer(
|
|
|
self.block,
|
|
|
self.inplanes,
|
|
|
planes,
|
|
|
num_blocks,
|
|
|
stride=stride,
|
|
|
dilation=dilation,
|
|
|
factorize=self.stage_factorization[i],
|
|
|
norm_cfg=self.norm_cfg,
|
|
|
with_cp=with_cp)
|
|
|
self.inplanes = planes * self.block.expansion
|
|
|
layer_name = f'layer{i + 1}'
|
|
|
self.add_module(layer_name, res_layer)
|
|
|
self.res_layers.append(layer_name)
|
|
|
|
|
|
self.feat_dim = self.block.expansion * self.base_channels * 2**(
|
|
|
len(self.stage_blocks) - 1)
|
|
|
|
|
|
@staticmethod
|
|
|
def make_res_layer(block: nn.Module,
|
|
|
inplanes: int,
|
|
|
planes: int,
|
|
|
blocks: int,
|
|
|
stride: int = 1,
|
|
|
dilation: int = 1,
|
|
|
factorize: int = 1,
|
|
|
norm_cfg: Optional[ConfigType] = None,
|
|
|
with_cp: bool = False) -> nn.Module:
|
|
|
"""Build residual layer for ResNetAudio.
|
|
|
|
|
|
Args:
|
|
|
block (nn.Module): Residual module to be built.
|
|
|
inplanes (int): Number of channels for the input feature
|
|
|
in each block.
|
|
|
planes (int): Number of channels for the output feature
|
|
|
in each block.
|
|
|
blocks (int): Number of residual blocks.
|
|
|
stride (int): Strides of residual blocks of each stage.
|
|
|
Defaults to 1.
|
|
|
dilation (int): Spacing between kernel elements. Defaults to 1.
|
|
|
factorize (Uninon[int, Sequence[int]]): Determine whether to
|
|
|
factorize for each block. Defaults to 1.
|
|
|
norm_cfg (Union[dict, ConfigDict], optional): Config for norm
|
|
|
layers. Defaults to None.
|
|
|
with_cp (bool): Use checkpoint or not. Using checkpoint will save
|
|
|
some memory while slowing down the training speed.
|
|
|
Defaults to False.
|
|
|
|
|
|
Returns:
|
|
|
nn.Module: A residual layer for the given config.
|
|
|
"""
|
|
|
factorize = factorize if not isinstance(
|
|
|
factorize, int) else (factorize, ) * blocks
|
|
|
assert len(factorize) == blocks
|
|
|
downsample = None
|
|
|
if stride != 1 or inplanes != planes * block.expansion:
|
|
|
downsample = ConvModule(
|
|
|
inplanes,
|
|
|
planes * block.expansion,
|
|
|
kernel_size=1,
|
|
|
stride=stride,
|
|
|
bias=False,
|
|
|
norm_cfg=norm_cfg,
|
|
|
act_cfg=None)
|
|
|
|
|
|
layers = []
|
|
|
layers.append(
|
|
|
block(
|
|
|
inplanes,
|
|
|
planes,
|
|
|
stride,
|
|
|
dilation,
|
|
|
downsample,
|
|
|
factorize=(factorize[0] == 1),
|
|
|
norm_cfg=norm_cfg,
|
|
|
with_cp=with_cp))
|
|
|
inplanes = planes * block.expansion
|
|
|
for i in range(1, blocks):
|
|
|
layers.append(
|
|
|
block(
|
|
|
inplanes,
|
|
|
planes,
|
|
|
1,
|
|
|
dilation,
|
|
|
factorize=(factorize[i] == 1),
|
|
|
norm_cfg=norm_cfg,
|
|
|
with_cp=with_cp))
|
|
|
|
|
|
return nn.Sequential(*layers)
|
|
|
|
|
|
def _make_stem_layer(self) -> None:
|
|
|
"""Construct the stem layers consists of a ``conv+norm+act`` module and
|
|
|
a pooling layer."""
|
|
|
self.conv1 = ConvModule(
|
|
|
self.in_channels,
|
|
|
self.base_channels,
|
|
|
kernel_size=self.conv1_kernel,
|
|
|
stride=self.conv1_stride,
|
|
|
bias=False,
|
|
|
conv_cfg=dict(type='ConvAudio', op='sum'),
|
|
|
norm_cfg=self.norm_cfg,
|
|
|
act_cfg=self.act_cfg)
|
|
|
|
|
|
def _freeze_stages(self) -> None:
|
|
|
"""Prevent all the parameters from being optimized before
|
|
|
``self.frozen_stages``."""
|
|
|
if self.frozen_stages >= 0:
|
|
|
self.conv1.bn.eval()
|
|
|
for m in [self.conv1.conv, self.conv1.bn]:
|
|
|
for param in m.parameters():
|
|
|
param.requires_grad = False
|
|
|
|
|
|
for i in range(1, self.frozen_stages + 1):
|
|
|
m = getattr(self, f'layer{i}')
|
|
|
m.eval()
|
|
|
for param in m.parameters():
|
|
|
param.requires_grad = False
|
|
|
|
|
|
def init_weights(self) -> None:
|
|
|
"""Initiate the parameters either from existing checkpoint or from
|
|
|
scratch."""
|
|
|
if isinstance(self.pretrained, str):
|
|
|
logger = MMLogger.get_current_instance()
|
|
|
logger.info(f'load model from: {self.pretrained}')
|
|
|
|
|
|
load_checkpoint(self, self.pretrained, strict=False, logger=logger)
|
|
|
|
|
|
elif self.pretrained is None:
|
|
|
for m in self.modules():
|
|
|
if isinstance(m, nn.Conv2d):
|
|
|
kaiming_init(m)
|
|
|
elif isinstance(m, _BatchNorm):
|
|
|
constant_init(m, 1)
|
|
|
|
|
|
if self.zero_init_residual:
|
|
|
for m in self.modules():
|
|
|
if isinstance(m, Bottleneck2dAudio):
|
|
|
constant_init(m.conv3.bn, 0)
|
|
|
|
|
|
else:
|
|
|
raise TypeError('pretrained must be a str or None')
|
|
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
|
"""Defines the computation performed at every call.
|
|
|
|
|
|
Args:
|
|
|
x (torch.Tensor): The input data.
|
|
|
|
|
|
Returns:
|
|
|
torch.Tensor: The feature of the input samples extracted
|
|
|
by the backbone.
|
|
|
"""
|
|
|
x = self.conv1(x)
|
|
|
for layer_name in self.res_layers:
|
|
|
res_layer = getattr(self, layer_name)
|
|
|
x = res_layer(x)
|
|
|
return x
|
|
|
|
|
|
def train(self, mode: bool = True) -> None:
|
|
|
"""Set the optimization status when training."""
|
|
|
super().train(mode)
|
|
|
self._freeze_stages()
|
|
|
if mode and self.norm_eval:
|
|
|
for m in self.modules():
|
|
|
if isinstance(m, _BatchNorm):
|
|
|
m.eval()
|
|
|
|
