def __init__(self,
inplanes,
planes,
stride=1,
downsample=None,
use_att=False,
att_mode='ours'):
super(Bottleneck, self).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(planes,
planes,
kernel_size=3,
stride=stride,
padding=1,
bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.conv3 = nn.Conv2d(planes,
planes * self.expansion,
kernel_size=1,
bias=False)
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
self.use_att = use_att
if use_att:
assert att_mode in ['ours', 'cbam', 'se']
if att_mode == 'ours':
self.att = AttentionModule(planes * self.expansion)
elif att_mode == 'cbam':
self.att = CBAM(planes * self.expansion)
elif att_mode == 'se':
self.att = SELayer(planes * self.expansion)
def __init__(self,
inplanes,
planes,
stride=1,
downsample=None,
use_att=False,
att_mode='ours'):
super(BasicBlock, self).__init__()
self.conv1 = conv3x3(inplanes, planes, stride)
self.bn1 = nn.BatchNorm2d(planes)
self.relu = nn.ReLU(inplace=True)
self.conv2 = conv3x3(planes, planes)
self.bn2 = nn.BatchNorm2d(planes)
self.downsample = downsample
self.stride = stride
self.use_att = use_att
if use_att:
assert att_mode in ['ours', 'cbam', 'se']
if att_mode == 'ours':
self.att = AttentionModule(planes)
elif att_mode == 'cbam':
self.att = CBAM(planes)
elif att_mode == 'se':
self.att = SELayer(planes)
def __init__(self,
num_class,
num_segments,
modality,
base_model,
new_length=None,
consensus_type='avg',
dropout=0.5):
super(TSN, self).__init__()
self.modality = modality
self.num_segments = num_segments
self.dropout = dropout
self.base_model = base_model
self.new_length = new_length
self.consensus_type = consensus_type
print(("""
Initializing TSN with base model: P3D.
TSN Configurations:
input_modality: {}
num_segments: {}
new_length: {}
consensus_module: {}
dropout_ratio: {}
""".format(self.modality, self.num_segments, self.new_length,
consensus_type, self.dropout)))
self.attention = AttentionModule()
self.avgpool = nn.AvgPool3d(kernel_size=(self.num_segments, 5, 5),
stride=1)
self.dropout = nn.Dropout(p=dropout)
self.fc = nn.Linear(2048, 101)
def __init__(self,
block,
layers,
num_classes=1000,
use_att=False,
att_mode='ours'):
self.inplanes = 64
super(ResNet, self).__init__()
self.conv1 = nn.Conv2d(3,
64,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block,
64,
layers[0],
stride=1,
use_att=False,
att_mode=att_mode)
self.layer2 = self._make_layer(block,
128,
layers[1],
stride=2,
use_att=False,
att_mode=att_mode)
self.layer3 = self._make_layer(block,
256,
layers[2],
stride=2,
use_att=False,
att_mode=att_mode)
self.layer4 = self._make_layer(block,
512,
layers[3],
stride=2,
use_att=False,
att_mode=att_mode)
self.avgpool = nn.AvgPool2d(kernel_size=7, stride=1)
self.final_fc = nn.Linear(512 * block.expansion, num_classes)
self.use_att = use_att
if use_att:
self.att = AttentionModule(512 * block.expansion)
self.num_classes = num_classes
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)