|
|
| from copy import deepcopy
|
| from typing import Optional
|
|
|
| import torch
|
| import torch.nn as nn
|
| from torch.utils import checkpoint as cp
|
|
|
| from mmaction.registry import MODELS
|
| from ..common import TAM
|
| from .resnet import Bottleneck, ResNet
|
|
|
|
|
| class TABlock(nn.Module):
|
| """Temporal Adaptive Block (TA-Block) for TANet.
|
|
|
| This block is proposed in `TAM: TEMPORAL ADAPTIVE MODULE FOR VIDEO
|
| RECOGNITION <https://arxiv.org/pdf/2005.06803>`_
|
|
|
| The temporal adaptive module (TAM) is embedded into ResNet-Block
|
| after the first Conv2D, which turns the vanilla ResNet-Block
|
| into TA-Block.
|
|
|
| Args:
|
| block (nn.Module): Residual blocks to be substituted.
|
| num_segments (int): Number of frame segments.
|
| tam_cfg (dict): Config for temporal adaptive module (TAM).
|
| """
|
|
|
| def __init__(self, block: nn.Module, num_segments: int,
|
| tam_cfg: dict) -> None:
|
| super().__init__()
|
| self.tam_cfg = deepcopy(tam_cfg)
|
| self.block = block
|
| self.num_segments = num_segments
|
| self.tam = TAM(
|
| in_channels=block.conv1.out_channels,
|
| num_segments=num_segments,
|
| **self.tam_cfg)
|
|
|
| if not isinstance(self.block, Bottleneck):
|
| raise NotImplementedError('TA-Blocks have not been fully '
|
| 'implemented except the pattern based '
|
| 'on Bottleneck block.')
|
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor:
|
| """Defines the computation performed at every call."""
|
| assert isinstance(self.block, Bottleneck)
|
|
|
| def _inner_forward(x):
|
| """Forward wrapper for utilizing checkpoint."""
|
| identity = x
|
|
|
| out = self.block.conv1(x)
|
| out = self.tam(out)
|
| out = self.block.conv2(out)
|
| out = self.block.conv3(out)
|
|
|
| if self.block.downsample is not None:
|
| identity = self.block.downsample(x)
|
|
|
| out = out + identity
|
|
|
| return out
|
|
|
| if self.block.with_cp and x.requires_grad:
|
| out = cp.checkpoint(_inner_forward, x)
|
| else:
|
| out = _inner_forward(x)
|
|
|
| out = self.block.relu(out)
|
|
|
| return out
|
|
|
|
|
| @MODELS.register_module()
|
| class TANet(ResNet):
|
| """Temporal Adaptive Network (TANet) backbone.
|
|
|
| This backbone is proposed in `TAM: TEMPORAL ADAPTIVE MODULE FOR VIDEO
|
| RECOGNITION <https://arxiv.org/pdf/2005.06803>`_
|
|
|
| Embedding the temporal adaptive module (TAM) into ResNet to
|
| instantiate TANet.
|
|
|
| Args:
|
| depth (int): Depth of resnet, from ``{18, 34, 50, 101, 152}``.
|
| num_segments (int): Number of frame segments.
|
| tam_cfg (dict, optional): Config for temporal adaptive module (TAM).
|
| Defaults to None.
|
| """
|
|
|
| def __init__(self,
|
| depth: int,
|
| num_segments: int,
|
| tam_cfg: Optional[dict] = None,
|
| **kwargs) -> None:
|
| super().__init__(depth, **kwargs)
|
| assert num_segments >= 3
|
| self.num_segments = num_segments
|
| tam_cfg = dict() if tam_cfg is None else tam_cfg
|
| self.tam_cfg = deepcopy(tam_cfg)
|
| super().init_weights()
|
| self.make_tam_modeling()
|
|
|
| def init_weights(self):
|
| """Initialize weights."""
|
| pass
|
|
|
| def make_tam_modeling(self):
|
| """Replace ResNet-Block with TA-Block."""
|
|
|
| def make_tam_block(stage, num_segments, tam_cfg=dict()):
|
| blocks = list(stage.children())
|
| for i, block in enumerate(blocks):
|
| blocks[i] = TABlock(block, num_segments, deepcopy(tam_cfg))
|
| return nn.Sequential(*blocks)
|
|
|
| for i in range(self.num_stages):
|
| layer_name = f'layer{i + 1}'
|
| res_layer = getattr(self, layer_name)
|
| setattr(self, layer_name,
|
| make_tam_block(res_layer, self.num_segments, self.tam_cfg))
|
|
|
