|
|
| from typing import Tuple, Union
|
|
|
| import torch
|
| import torch.nn as nn
|
| from mmcv.cnn import ConvModule
|
| from mmengine.model.weight_init import constant_init, kaiming_init
|
| from torch.nn.modules.utils import _pair
|
|
|
| from mmaction.registry import MODELS
|
|
|
|
|
| @MODELS.register_module()
|
| class ConvAudio(nn.Module):
|
| """Conv2d module for AudioResNet backbone.
|
|
|
| <https://arxiv.org/abs/2001.08740>`_.
|
|
|
| Args:
|
| in_channels (int): Same as ``nn.Conv2d``.
|
| out_channels (int): Same as ``nn.Conv2d``.
|
| kernel_size (Union[int, Tuple[int]]): Same as ``nn.Conv2d``.
|
| op (str): Operation to merge the output of freq
|
| and time feature map. Choices are ``sum`` and ``concat``.
|
| Defaults to ``concat``.
|
| stride (Union[int, Tuple[int]]): Same as ``nn.Conv2d``. Defaults to 1.
|
| padding (Union[int, Tuple[int]]): Same as ``nn.Conv2d``. Defaults to 0.
|
| dilation (Union[int, Tuple[int]]): Same as ``nn.Conv2d``.
|
| Defaults to 1.
|
| groups (int): Same as ``nn.Conv2d``. Defaults to 1.
|
| bias (Union[bool, str]): If specified as ``auto``, it will be decided
|
| by the ``norm_cfg``. Bias will be set as True if ``norm_cfg``
|
| is None, otherwise False. Defaults to False.
|
| """
|
|
|
| def __init__(self,
|
| in_channels: int,
|
| out_channels: int,
|
| kernel_size: Union[int, Tuple[int]],
|
| op: str = 'concat',
|
| stride: Union[int, Tuple[int]] = 1,
|
| padding: Union[int, Tuple[int]] = 0,
|
| dilation: Union[int, Tuple[int]] = 1,
|
| groups: int = 1,
|
| bias: Union[bool, str] = False) -> None:
|
| super().__init__()
|
|
|
| kernel_size = _pair(kernel_size)
|
| stride = _pair(stride)
|
| padding = _pair(padding)
|
|
|
| self.in_channels = in_channels
|
| self.out_channels = out_channels
|
| self.kernel_size = kernel_size
|
| assert op in ['concat', 'sum']
|
| self.op = op
|
| self.stride = stride
|
| self.padding = padding
|
| self.dilation = dilation
|
| self.groups = groups
|
| self.bias = bias
|
| self.output_padding = (0, 0)
|
| self.transposed = False
|
|
|
| self.conv_1 = ConvModule(
|
| in_channels,
|
| out_channels,
|
| kernel_size=(kernel_size[0], 1),
|
| stride=stride,
|
| padding=(kernel_size[0] // 2, 0),
|
| bias=bias,
|
| conv_cfg=dict(type='Conv'),
|
| norm_cfg=dict(type='BN'),
|
| act_cfg=dict(type='ReLU'))
|
|
|
| self.conv_2 = ConvModule(
|
| in_channels,
|
| out_channels,
|
| kernel_size=(1, kernel_size[1]),
|
| stride=stride,
|
| padding=(0, kernel_size[1] // 2),
|
| bias=bias,
|
| conv_cfg=dict(type='Conv'),
|
| norm_cfg=dict(type='BN'),
|
| act_cfg=dict(type='ReLU'))
|
|
|
| self.init_weights()
|
|
|
| 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.
|
| """
|
| x_1 = self.conv_1(x)
|
| x_2 = self.conv_2(x)
|
| if self.op == 'concat':
|
| out = torch.cat([x_1, x_2], 1)
|
| else:
|
| out = x_1 + x_2
|
| return out
|
|
|
| def init_weights(self) -> None:
|
| """Initiate the parameters from scratch."""
|
| kaiming_init(self.conv_1.conv)
|
| kaiming_init(self.conv_2.conv)
|
| constant_init(self.conv_1.bn, 1, bias=0)
|
| constant_init(self.conv_2.bn, 1, bias=0)
|
|
|
