384 lines
12 KiB
Python
384 lines
12 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from detectron2.modeling import Backbone
|
|
|
|
|
|
def conv1x3x3(in_planes, out_planes, stride=(1, 1, 1)):
|
|
return nn.Conv3d(
|
|
in_planes,
|
|
out_planes,
|
|
kernel_size=(1, 3, 3),
|
|
stride=stride,
|
|
padding=(0, 1, 1),
|
|
bias=False)
|
|
|
|
|
|
def conv3x3x3(in_planes, out_planes, stride=(1, 1, 1)):
|
|
return nn.Conv3d(
|
|
in_planes,
|
|
out_planes,
|
|
kernel_size=(3, 3, 3),
|
|
stride=stride,
|
|
padding=(1, 1, 1),
|
|
bias=False)
|
|
|
|
|
|
def conv1x1x1(in_planes, out_planes, stride=(1, 1, 1)):
|
|
return nn.Conv3d(
|
|
in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
|
|
|
|
|
def conv3x1x1(in_planes, out_planes, stride=(1, 1, 1)):
|
|
return nn.Conv3d(
|
|
in_planes,
|
|
out_planes,
|
|
kernel_size=(3, 1, 1),
|
|
stride=stride,
|
|
padding=(1, 0, 0),
|
|
bias=False)
|
|
|
|
|
|
class BasicBlock3D(nn.Module):
|
|
expansion = 1
|
|
|
|
def __init__(self,
|
|
inplanes,
|
|
planes,
|
|
stride=1,
|
|
op='c2d',
|
|
downsample=None,
|
|
base_width=64,
|
|
norm_layer=None):
|
|
super(BasicBlock3D, self).__init__()
|
|
dilation = 1
|
|
groups = 1
|
|
if norm_layer is None:
|
|
norm_layer = nn.BatchNorm3d
|
|
if groups != 1 or base_width != 64:
|
|
raise ValueError(
|
|
'BasicBlock only supports groups=1 and base_width=64')
|
|
if dilation > 1:
|
|
raise NotImplementedError(
|
|
'Dilation > 1 not supported in BasicBlock')
|
|
stride = [stride] * 3 if isinstance(stride, int) else stride
|
|
self.t_stride = stride[0]
|
|
self.stride = stride
|
|
stride = [1] + list(stride[1:])
|
|
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
|
|
# self.conv1 = conv3x3(inplanes, planes, stride)
|
|
self.conv1 = conv1x3x3(inplanes, planes, stride)
|
|
self.bn1 = norm_layer(planes)
|
|
self.relu = nn.ReLU(inplace=True)
|
|
self.conv1d = conv3x1x1(planes, planes)
|
|
# self.conv2 = conv3x3(planes, planes)
|
|
self.conv2 = conv1x3x3(planes, planes)
|
|
self.bn2 = norm_layer(planes)
|
|
self.downsample = downsample
|
|
|
|
def forward(self, x):
|
|
identity = x
|
|
|
|
out = self.conv1(x)
|
|
if self.t_stride > 1:
|
|
out = torch.max_pool3d(
|
|
out, [self.t_stride, 1, 1], stride=[self.t_stride, 1, 1])
|
|
out = self.bn1(out)
|
|
out = self.relu(out)
|
|
|
|
out = self.conv1d(out)
|
|
|
|
out = self.conv2(out)
|
|
out = self.bn2(out)
|
|
|
|
if self.downsample is not None:
|
|
identity = self.downsample(x)
|
|
|
|
out += identity
|
|
out = self.relu(out)
|
|
|
|
return out
|
|
|
|
@property
|
|
def out_channels(self):
|
|
return self.conv2.out_channels
|
|
|
|
|
|
class Bottleneck3D(nn.Module):
|
|
expansion = 2
|
|
|
|
def __init__(self,
|
|
inplanes,
|
|
planes,
|
|
stride=1,
|
|
op='c2d',
|
|
downsample=None,
|
|
base_width=64,
|
|
norm_layer=None):
|
|
super(Bottleneck3D, self).__init__()
|
|
self.op = op
|
|
if norm_layer is None:
|
|
norm_layer = nn.BatchNorm3d
|
|
width = int(planes * (base_width / 64.))
|
|
stride = [stride] * 3 if isinstance(stride, int) else stride
|
|
self.conv1 = conv3x1x1(inplanes, width) if op == 'p3d' else conv1x1x1(
|
|
inplanes, width)
|
|
self.bn1 = norm_layer(width)
|
|
self.conv2 = conv3x3x3(width, width,
|
|
stride) if op == 'c3d' else conv1x3x3(
|
|
width, width, stride)
|
|
self.bn2 = norm_layer(width)
|
|
self.conv3 = conv1x1x1(width, planes * self.expansion)
|
|
self.bn3 = norm_layer(planes * self.expansion)
|
|
self.relu = nn.ReLU(inplace=True)
|
|
self.downsample = downsample
|
|
|
|
def forward(self, x):
|
|
identity = x
|
|
|
|
out = self.conv1(x)
|
|
out = self.bn1(out)
|
|
out = self.relu(out)
|
|
|
|
out = self.conv2(out)
|
|
out = self.bn2(out)
|
|
out = self.relu(out)
|
|
|
|
if self.op == 'tsm':
|
|
out = self.tsm(out)
|
|
|
|
out = self.conv3(out)
|
|
out = self.bn3(out)
|
|
|
|
if self.downsample is not None:
|
|
identity = self.downsample(x)
|
|
|
|
out += identity
|
|
out = self.relu(out)
|
|
|
|
return out
|
|
|
|
@property
|
|
def out_channels(self):
|
|
return self.conv3.out_channels
|
|
|
|
|
|
class ResNet3D(Backbone):
|
|
|
|
def __init__(self,
|
|
block,
|
|
layers,
|
|
ops,
|
|
t_stride,
|
|
num_classes=1000,
|
|
zero_init_residual=False,
|
|
groups=1,
|
|
width_per_group=64,
|
|
replace_stride_with_dilation=None,
|
|
reduce_dim=0,
|
|
norm_layer=None):
|
|
self.reduce_dim = reduce_dim
|
|
self.num_classes = num_classes
|
|
super(ResNet3D, self).__init__()
|
|
if norm_layer is None:
|
|
norm_layer = nn.BatchNorm3d
|
|
self._norm_layer = norm_layer
|
|
self._out_feature_strides = {'res3': 8, 'res4': 16, 'res5': 32}
|
|
self._out_features = ['res3', 'res4', 'res5']
|
|
self._out_feature_channels = {}
|
|
self.outputs = {}
|
|
self.inplanes = 64
|
|
if replace_stride_with_dilation is None:
|
|
# each element in the tuple indicates if we should replace
|
|
# the 2x2 stride with a dilated convolution instead
|
|
replace_stride_with_dilation = [False, False, False]
|
|
if len(replace_stride_with_dilation) != 3:
|
|
raise ValueError('replace_stride_with_dilation should be None '
|
|
'or a 3-element tuple, got {}'.format(
|
|
replace_stride_with_dilation))
|
|
self.groups = groups
|
|
self.base_width = width_per_group
|
|
self.conv1 = nn.Conv3d(
|
|
3,
|
|
self.inplanes, (1, 7, 7),
|
|
stride=(t_stride[0], 2, 2),
|
|
padding=(0, 3, 3),
|
|
bias=False)
|
|
self.bn1 = norm_layer(self.inplanes)
|
|
self.relu = nn.ReLU(inplace=True)
|
|
self.maxpool = nn.MaxPool3d(
|
|
kernel_size=(1, 3, 3),
|
|
stride=(t_stride[1], 2, 2),
|
|
padding=(0, 1, 1))
|
|
self.layer1 = self._make_layer(
|
|
block, 64, layers[0], stride=(1, 1, 1), ops=ops[:layers[0]])
|
|
self.layer2 = self._make_layer(
|
|
block,
|
|
128,
|
|
layers[1],
|
|
stride=(t_stride[2], 2, 2),
|
|
ops=ops[layers[0]:][:layers[1]])
|
|
self.layer3 = self._make_layer(
|
|
block,
|
|
256,
|
|
layers[2],
|
|
stride=(t_stride[3], 2, 2),
|
|
ops=ops[sum(layers[:2], 0):][:layers[2]])
|
|
self.layer4 = self._make_layer(
|
|
block,
|
|
512,
|
|
layers[3],
|
|
stride=(t_stride[4], 2, 2),
|
|
ops=ops[sum(layers[:3], 0):][:layers[3]])
|
|
if num_classes is not None:
|
|
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
|
|
self.sptial_atten = nn.Conv2d(2, 1, kernel_size=7, padding=3)
|
|
self.drop = nn.Dropout(0.5)
|
|
if reduce_dim > 0:
|
|
self.rd_conv = nn.Conv2d(
|
|
512 * block.expansion, reduce_dim, kernel_size=1)
|
|
self.clc = nn.Conv2d(reduce_dim, num_classes, kernel_size=1)
|
|
else:
|
|
self.clc = nn.Conv2d(
|
|
512 * block.expansion, num_classes, kernel_size=1)
|
|
|
|
self._out_feature_channels['res3'] = self.layer2[-1].out_channels
|
|
self._out_feature_channels['res4'] = self.layer3[-1].out_channels
|
|
self._out_feature_channels['res5'] = self.layer4[-1].out_channels
|
|
|
|
for m in self.modules():
|
|
if isinstance(m, nn.Conv3d):
|
|
nn.init.kaiming_normal_(
|
|
m.weight, mode='fan_out', nonlinearity='relu')
|
|
elif isinstance(m, (nn.BatchNorm3d, nn.GroupNorm)):
|
|
nn.init.constant_(m.weight, 1)
|
|
nn.init.constant_(m.bias, 0)
|
|
|
|
# Zero-initialize the last BN in each residual branch,
|
|
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
|
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
|
if zero_init_residual:
|
|
for m in self.modules():
|
|
if isinstance(m, Bottleneck3D):
|
|
nn.init.constant_(m.bn3.weight, 0)
|
|
|
|
def _make_layer(self, block, planes, blocks, ops, stride=(1, 1, 1)):
|
|
norm_layer = self._norm_layer
|
|
downsample = None
|
|
if stride != (1, 1, 1) or self.inplanes != planes * block.expansion:
|
|
downsample = nn.Sequential(
|
|
conv1x1x1(self.inplanes, planes * block.expansion, stride),
|
|
norm_layer(planes * block.expansion),
|
|
)
|
|
|
|
layers = []
|
|
layers.append(
|
|
block(self.inplanes, planes, stride, ops[0], downsample,
|
|
self.base_width, norm_layer))
|
|
self.inplanes = planes * block.expansion
|
|
for i in range(1, blocks):
|
|
layers.append(
|
|
block(
|
|
self.inplanes,
|
|
planes,
|
|
op=ops[i],
|
|
base_width=self.base_width,
|
|
norm_layer=norm_layer))
|
|
|
|
return nn.Sequential(*layers)
|
|
|
|
def features(self, x):
|
|
x = self.norm_x(x)
|
|
x = self.conv1(x)
|
|
x = self.bn1(x)
|
|
x = self.relu(x)
|
|
x = self.maxpool(x)
|
|
|
|
x = self.layer1(x)
|
|
x = self.layer2(x)
|
|
self.outputs['res3'] = x.mean(dim=2)
|
|
x = self.layer3(x)
|
|
self.outputs['res4'] = x.mean(dim=2)
|
|
x = self.layer4(x) # N,C,T,H,W
|
|
self.outputs['res5'] = x.mean(dim=2)
|
|
if self.num_classes is not None:
|
|
x = torch.mean(x, dim=2) # 解决时间维度, N,C,H,W
|
|
# spatial attention
|
|
ftr = torch.cat(
|
|
(x.max(dim=1, keepdim=True)[0], x.mean(dim=1, keepdim=True)),
|
|
dim=1)
|
|
score = self.sptial_atten(ftr) # N,1,H,W
|
|
x = x * torch.sigmoid(score) # N,C,H,W
|
|
self.score = score
|
|
|
|
x = self.avgpool(x) # ,N,C,1,1
|
|
if self.reduce_dim > 0:
|
|
x = self.rd_conv(x)
|
|
self.outputs['ftr'] = x.mean(dim=(2, 3))
|
|
return x
|
|
|
|
def logits(self, x):
|
|
x = self.features(x)
|
|
x = self.clc(x)
|
|
return x
|
|
|
|
def forward(self, x):
|
|
ftr = self.features(x)
|
|
if self.num_classes is not None:
|
|
x = self.drop(ftr)
|
|
x = self.clc(x)
|
|
x = torch.mean(x, (2, 3))
|
|
return x
|
|
|
|
return self.outputs
|
|
|
|
@torch.no_grad()
|
|
def norm_x(self, x):
|
|
m = x.new_tensor([0.485, 0.456, 0.406]).reshape([1, 3, 1, 1, 1])
|
|
s = x.new_tensor([0.229, 0.224, 0.225]).reshape([1, 3, 1, 1, 1])
|
|
x -= m
|
|
x /= s
|
|
return x
|
|
|
|
|
|
def resnet101_3d(ops, t_stride, num_class, reduce_dim=0):
|
|
net = ResNet3D(
|
|
Bottleneck3D, [3, 4, 23, 3],
|
|
ops=ops,
|
|
t_stride=t_stride,
|
|
num_classes=num_class,
|
|
reduce_dim=reduce_dim)
|
|
return net
|
|
|
|
|
|
def resnet50_3d(ops, t_stride, num_class, reduce_dim=0):
|
|
net = ResNet3D(
|
|
Bottleneck3D, [3, 4, 6, 3],
|
|
ops=ops,
|
|
t_stride=t_stride,
|
|
num_classes=num_class,
|
|
reduce_dim=reduce_dim)
|
|
return net
|
|
|
|
|
|
def resnet34_3d(ops, t_stride, num_class, reduce_dim=0):
|
|
net = ResNet3D(
|
|
BasicBlock3D, [3, 4, 6, 3],
|
|
ops=ops,
|
|
t_stride=t_stride,
|
|
num_classes=num_class,
|
|
reduce_dim=reduce_dim)
|
|
return net
|
|
|
|
|
|
def resnet18_3d(ops, t_stride, num_class, reduce_dim=0):
|
|
net = ResNet3D(
|
|
BasicBlock3D, [2, 2, 2, 2],
|
|
ops=ops,
|
|
t_stride=t_stride,
|
|
num_classes=num_class,
|
|
reduce_dim=reduce_dim)
|
|
return net
|