|
|
|
|
|
from typing import Tuple, Union
|
|
|
|
|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
from mmcv.cnn import ConvModule
|
|
|
|
|
|
from mmaction.models.backbones.resnet import ResNet
|
|
|
from mmaction.registry import MODELS
|
|
|
|
|
|
|
|
|
@MODELS.register_module()
|
|
|
class C2D(ResNet):
|
|
|
"""C2D backbone.
|
|
|
|
|
|
Compared to ResNet-50, a temporal-pool is added after the first
|
|
|
bottleneck. Detailed structure is kept same as "video-nonlocal-net" repo.
|
|
|
Please refer to https://github.com/facebookresearch/video-nonlocal-net/blob
|
|
|
/main/scripts/run_c2d_baseline_400k.sh.
|
|
|
Please note that there are some improvements compared to "Non-local Neural
|
|
|
Networks" paper (https://arxiv.org/abs/1711.07971).
|
|
|
Differences are noted at https://github.com/facebookresearch/video-nonlocal
|
|
|
-net#modifications-for-improving-speed.
|
|
|
"""
|
|
|
|
|
|
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,
|
|
|
64,
|
|
|
kernel_size=7,
|
|
|
stride=2,
|
|
|
padding=3,
|
|
|
bias=False,
|
|
|
conv_cfg=self.conv_cfg,
|
|
|
norm_cfg=self.norm_cfg,
|
|
|
act_cfg=self.act_cfg)
|
|
|
self.maxpool3d_1 = nn.MaxPool3d(
|
|
|
kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 0, 0))
|
|
|
self.maxpool3d_2 = nn.MaxPool3d(
|
|
|
kernel_size=(2, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
|
|
|
|
|
|
def forward(self, x: torch.Tensor) \
|
|
|
-> Union[torch.Tensor, Tuple[torch.Tensor]]:
|
|
|
"""Defines the computation performed at every call.
|
|
|
|
|
|
Args:
|
|
|
x (torch.Tensor): The input data.
|
|
|
|
|
|
Returns:
|
|
|
Union[torch.Tensor or Tuple[torch.Tensor]]: The feature of the
|
|
|
input samples extracted by the backbone.
|
|
|
"""
|
|
|
|
|
|
batches = x.shape[0]
|
|
|
|
|
|
def _convert_to_2d(x: torch.Tensor) -> torch.Tensor:
|
|
|
"""(N, C, T, H, W) -> (N x T, C, H, W)"""
|
|
|
x = x.permute((0, 2, 1, 3, 4))
|
|
|
x = x.reshape(-1, x.shape[2], x.shape[3], x.shape[4])
|
|
|
return x
|
|
|
|
|
|
def _convert_to_3d(x: torch.Tensor) -> torch.Tensor:
|
|
|
"""(N x T, C, H, W) -> (N, C, T, H, W)"""
|
|
|
x = x.reshape(batches, -1, x.shape[1], x.shape[2], x.shape[3])
|
|
|
x = x.permute((0, 2, 1, 3, 4))
|
|
|
return x
|
|
|
|
|
|
x = _convert_to_2d(x)
|
|
|
x = self.conv1(x)
|
|
|
x = _convert_to_3d(x)
|
|
|
x = self.maxpool3d_1(x)
|
|
|
x = _convert_to_2d(x)
|
|
|
outs = []
|
|
|
for i, layer_name in enumerate(self.res_layers):
|
|
|
res_layer = getattr(self, layer_name)
|
|
|
x = res_layer(x)
|
|
|
if i == 0:
|
|
|
x = _convert_to_3d(x)
|
|
|
x = self.maxpool3d_2(x)
|
|
|
x = _convert_to_2d(x)
|
|
|
if i in self.out_indices:
|
|
|
x = _convert_to_3d(x)
|
|
|
outs.append(x)
|
|
|
if len(outs) == 1:
|
|
|
return outs[0]
|
|
|
|
|
|
return tuple(outs)
|
|
|
|
