335 lines
12 KiB
Python
335 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import collections
|
|
import math
|
|
from itertools import repeat
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from scepter.modules.model.backbone.video.bricks.base_branch import BaseBranch
|
|
from scepter.modules.model.registry import BRICKS
|
|
from scepter.modules.utils.config import Config, dict_to_yaml
|
|
|
|
|
|
def _ntuple(n):
|
|
def parse(x):
|
|
if isinstance(x, collections.abc.Iterable):
|
|
return tuple(x)
|
|
return tuple(repeat(x, n))
|
|
|
|
return parse
|
|
|
|
|
|
_single = _ntuple(1)
|
|
_pair = _ntuple(2)
|
|
_triple = _ntuple(3)
|
|
_quadruple = _ntuple(4)
|
|
|
|
|
|
class TAdaConv2d(nn.Module):
|
|
""" Performs temporally adaptive 2D convolution.
|
|
Currently, only application on 5D tensors is supported, which makes TAdaConv2d
|
|
essentially a 3D convolution with temporal kernel size of 1.
|
|
"""
|
|
def __init__(self,
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size,
|
|
stride=1,
|
|
padding=0,
|
|
dilation=1,
|
|
groups=1,
|
|
bias=True):
|
|
"""
|
|
Args:
|
|
in_channels (int): number of input channels.
|
|
out_channels (int): number of output channels.
|
|
kernel_size (Tuple[int]): kernel size of TAdaConv2d.
|
|
stride (int, Tuple[int]): stride for the convolution in TAdaConv2d.
|
|
padding (int, Tuple[int]): padding for the convolution in TAdaConv2d.
|
|
dilation (Tuple[int]): dilation of the convolution in TAdaConv2d.
|
|
groups (int): number of groups for TAdaConv2d.
|
|
bias (bool): whether to use bias in TAdaConv2d.
|
|
"""
|
|
super(TAdaConv2d, self).__init__()
|
|
|
|
kernel_size = _triple(kernel_size)
|
|
stride = _triple(stride)
|
|
padding = _triple(padding)
|
|
dilation = _triple(dilation)
|
|
|
|
assert kernel_size[0] == 1
|
|
assert stride[0] == 1
|
|
assert padding[0] == 0
|
|
assert dilation[0] == 1
|
|
|
|
self.in_channels = in_channels
|
|
self.out_channels = out_channels
|
|
self.kernel_size = kernel_size
|
|
self.stride = stride
|
|
self.padding = padding
|
|
self.dilation = dilation
|
|
self.groups = groups
|
|
|
|
# base weights (W_b)
|
|
self.weight = nn.Parameter(
|
|
torch.Tensor(1, 1, out_channels, in_channels // groups,
|
|
kernel_size[1], kernel_size[2]))
|
|
if bias:
|
|
self.bias = nn.Parameter(torch.Tensor(out_channels))
|
|
else:
|
|
self.register_parameter('bias', None)
|
|
|
|
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
|
if self.bias is not None:
|
|
fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
|
|
bound = 1 / math.sqrt(fan_in)
|
|
nn.init.uniform_(self.bias, -bound, bound)
|
|
|
|
def forward(self, x, alpha):
|
|
"""
|
|
Args:
|
|
x (tensor): feature to perform convolution on.
|
|
alpha (tensor): calibration weight for the base weights.
|
|
W_t = alpha_t * W_b
|
|
"""
|
|
_, _, c_out, c_in, kh, kw = self.weight.size()
|
|
b, c_in, t, h, w = x.size()
|
|
x = x.permute(0, 2, 1, 3, 4).reshape(1, -1, h, w)
|
|
|
|
# alpha: B, C, T, H(1), W(1) -> B, T, C, H(1), W(1) -> B, T, 1, C, H(1), W(1)
|
|
# corresponding to calibrating the input channel
|
|
weight = (alpha.permute(0, 2, 1, 3, 4).unsqueeze(2) *
|
|
self.weight).reshape(-1, c_in, kh, kw)
|
|
|
|
bias = None
|
|
if self.bias is not None:
|
|
raise NotImplementedError
|
|
else:
|
|
output = F.conv2d(x,
|
|
weight=weight,
|
|
bias=bias,
|
|
stride=self.stride[1:],
|
|
padding=self.padding[1:],
|
|
dilation=self.dilation[1:],
|
|
groups=self.groups * b * t)
|
|
|
|
output = output.view(b, t, c_out, output.size(-2),
|
|
output.size(-1)).permute(0, 2, 1, 3, 4)
|
|
|
|
return output
|
|
|
|
def extra_repr(self) -> str:
|
|
return f'{self.in_channels}, {self.out_channels}, ' \
|
|
f'kernel_size: {self.kernel_size}, stride: {self.stride}, ' \
|
|
f'padding: {self.padding}, dilation: {self.dilation}, ' \
|
|
f'groups: {self.groups}'
|
|
|
|
|
|
class RouteFuncMLP(nn.Module):
|
|
""" The routing function for generating the calibration weights.
|
|
"""
|
|
def __init__(self, c_in, ratio, kernels, bn_eps=1e-5, bn_mmt=0.1):
|
|
"""
|
|
Args:
|
|
c_in (int): number of input channels.
|
|
ratio (int): reduction ratio for the routing function.
|
|
kernels (list): temporal kernel size of the stacked 1D convolutions
|
|
"""
|
|
super(RouteFuncMLP, self).__init__()
|
|
self.c_in = c_in
|
|
self.avgpool = nn.AdaptiveAvgPool3d((None, 1, 1))
|
|
self.globalpool = nn.AdaptiveAvgPool3d(1)
|
|
self.g = nn.Conv3d(
|
|
in_channels=c_in,
|
|
out_channels=c_in,
|
|
kernel_size=(1, 1, 1),
|
|
padding=0,
|
|
)
|
|
self.a = nn.Conv3d(
|
|
in_channels=c_in,
|
|
out_channels=int(c_in // ratio),
|
|
kernel_size=(kernels[0], 1, 1),
|
|
padding=(kernels[0] // 2, 0, 0),
|
|
)
|
|
self.bn = nn.BatchNorm3d(int(c_in // ratio),
|
|
eps=bn_eps,
|
|
momentum=bn_mmt)
|
|
self.relu = nn.ReLU(inplace=True)
|
|
self.b = nn.Conv3d(in_channels=int(c_in // ratio),
|
|
out_channels=c_in,
|
|
kernel_size=(kernels[1], 1, 1),
|
|
padding=(kernels[1] // 2, 0, 0),
|
|
bias=False)
|
|
self.b.skip_init = True
|
|
self.b.weight.data.zero_() # to make sure the initial values
|
|
# for the output is 1.
|
|
|
|
def forward(self, x):
|
|
g = self.globalpool(x)
|
|
x = self.avgpool(x)
|
|
x = self.a(x + self.g(g))
|
|
x = self.bn(x)
|
|
x = self.relu(x)
|
|
x = self.b(x) + 1
|
|
return x
|
|
|
|
|
|
@BRICKS.register_class()
|
|
class TAdaConvBlockAvgPool(BaseBranch):
|
|
""" The TAdaConv branch with average pooling as the feature aggregation scheme.
|
|
|
|
For details, see
|
|
Ziyuan Huang, Shiwei Zhang, Liang Pan, Zhiwu Qing, Mingqian Tang, Ziwei Liu, and Marcelo H. Ang Jr.
|
|
"TAda! Temporally-Adaptive Convolutions for Video Understanding."
|
|
"""
|
|
para_dict = {
|
|
'DIM_IN': {
|
|
'value': 64,
|
|
'description': "the branch's dim in!"
|
|
},
|
|
'NUM_FILTERS': {
|
|
'value': 64,
|
|
'description': 'the num of filter!'
|
|
},
|
|
'DOWNSAMPLING': {
|
|
'value': True,
|
|
'description': 'downsample spatial data or not!'
|
|
},
|
|
'DOWNSAMPLING_TEMPORAL': {
|
|
'value': True,
|
|
'description': 'downsample temporal data or not!'
|
|
},
|
|
'EXPANISION_RATIO': {
|
|
'value': 2,
|
|
'description': 'expanision ratio for this branch!'
|
|
},
|
|
'ROUTE_FUNC_R': {
|
|
'value': 4,
|
|
'description': 'the route func r!'
|
|
},
|
|
'ROUTE_FUNC_K': {
|
|
'value': [3, 3],
|
|
'description': 'the route func k!'
|
|
},
|
|
'POOL_K': {
|
|
'value': [3, 1, 1],
|
|
'description': 'the pool k!'
|
|
},
|
|
'BN_PARAMS': {
|
|
'value':
|
|
None,
|
|
'description':
|
|
'bn params data, key/value align with torch.BatchNorm3d/2d/1d!'
|
|
}
|
|
}
|
|
para_dict.update(BaseBranch.para_dict)
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
self.dim_in = cfg.DIM_IN
|
|
self.num_filters = cfg.NUM_FILTERS
|
|
self.kernel_size = cfg.KERNEL_SIZE
|
|
self.downsampling = cfg.get('DOWNSAMPLING', True)
|
|
self.downsampling_temporal = cfg.get('DOWNSAMPLING_TEMPORAL', True)
|
|
self.expansion_ratio = cfg.get('EXPANISION_RATIO', 2)
|
|
self.route_func_r = cfg.get('ROUTE_FUNC_R', 4)
|
|
self.route_func_k = cfg.get('ROUTE_FUNC_K', [3, 3])
|
|
self.pool_k = cfg.get('POOL_K', [3, 1, 1])
|
|
# bn_params or {}
|
|
self.bn_params = cfg.get('BN_PARAMS', None) or dict()
|
|
if isinstance(self.bn_params, Config):
|
|
self.bn_params = self.bn_params.__dict__
|
|
if self.downsampling:
|
|
if self.downsampling_temporal:
|
|
self.stride = [2, 2, 2]
|
|
else:
|
|
self.stride = [1, 2, 2]
|
|
else:
|
|
self.stride = [1, 1, 1]
|
|
|
|
super(TAdaConvBlockAvgPool, self).__init__(cfg, logger=logger)
|
|
|
|
def _construct_simple_block(self):
|
|
raise NotImplementedError
|
|
|
|
def _construct_bottleneck(self):
|
|
self.a = nn.Conv3d(in_channels=self.dim_in,
|
|
out_channels=self.num_filters //
|
|
self.expansion_ratio,
|
|
kernel_size=(1, 1, 1),
|
|
stride=(1, 1, 1),
|
|
padding=0,
|
|
bias=False)
|
|
self.a_bn = nn.BatchNorm3d(self.num_filters // self.expansion_ratio,
|
|
**self.bn_params)
|
|
self.a_relu = nn.ReLU(inplace=True)
|
|
|
|
self.b = TAdaConv2d(
|
|
in_channels=self.num_filters // self.expansion_ratio,
|
|
out_channels=self.num_filters // self.expansion_ratio,
|
|
kernel_size=(1, self.kernel_size[1], self.kernel_size[2]),
|
|
stride=(1, self.stride[1], self.stride[2]),
|
|
padding=(0, self.kernel_size[1] // 2, self.kernel_size[2] // 2),
|
|
bias=False)
|
|
self.b_rf = RouteFuncMLP(c_in=self.num_filters // self.expansion_ratio,
|
|
ratio=self.route_func_r,
|
|
kernels=self.route_func_k)
|
|
self.b_bn = nn.BatchNorm3d(self.num_filters // self.expansion_ratio,
|
|
**self.bn_params)
|
|
self.b_avgpool = nn.AvgPool3d(kernel_size=self.pool_k,
|
|
stride=1,
|
|
padding=(self.pool_k[0] // 2,
|
|
self.pool_k[1] // 2,
|
|
self.pool_k[2] // 2))
|
|
self.b_avgpool_bn = nn.BatchNorm3d(
|
|
self.num_filters // self.expansion_ratio, **self.bn_params)
|
|
self.b_avgpool_bn.skip_init = True
|
|
self.b_avgpool_bn.weight.data.zero_()
|
|
self.b_avgpool_bn.bias.data.zero_()
|
|
|
|
self.b_relu = nn.ReLU(inplace=True)
|
|
|
|
self.c = nn.Conv3d(in_channels=self.num_filters //
|
|
self.expansion_ratio,
|
|
out_channels=self.num_filters,
|
|
kernel_size=(1, 1, 1),
|
|
stride=(1, 1, 1),
|
|
padding=0,
|
|
bias=False)
|
|
self.c_bn = nn.BatchNorm3d(self.num_filters, **self.bn_params)
|
|
|
|
def forward(self, x):
|
|
if self.branch_style == 'simple_block':
|
|
raise NotImplementedError
|
|
|
|
x = self.a(x)
|
|
x = self.a_bn(x)
|
|
x = self.a_relu(x)
|
|
x = self.b(x, self.b_rf(x))
|
|
x = self.b_bn(x) + self.b_avgpool_bn(self.b_avgpool(x))
|
|
x = self.b_relu(x)
|
|
|
|
x = self.c(x)
|
|
x = self.c_bn(x)
|
|
return x
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
'''
|
|
{ "ENV" :
|
|
{ "description" : "",
|
|
"A" : {
|
|
"value": 1.0,
|
|
"description": ""
|
|
}
|
|
}
|
|
}
|
|
:return:
|
|
'''
|
|
return dict_to_yaml('BRANCH',
|
|
__class__.__name__,
|
|
TAdaConvBlockAvgPool.para_dict,
|
|
set_name=True)
|