""" Simplified from torchvision.models.video.r3d_18. The citation information is shown below. @article{DBLP:journals/corr/abs-1711-11248, author = {Du Tran and Heng Wang and Lorenzo Torresani and Jamie Ray and Yann LeCun and Manohar Paluri}, title = {A Closer Look at Spatiotemporal Convolutions for Action Recognition}, journal = {CoRR}, volume = {abs/1711.11248}, year = {2017}, url = {http://arxiv.org/abs/1711.11248}, archivePrefix = {arXiv}, eprint = {1711.11248}, timestamp = {Mon, 13 Aug 2018 16:46:39 +0200}, biburl = {https://dblp.org/rec/journals/corr/abs-1711-11248.bib}, bibsource = {dblp computer science bibliography, https://dblp.org} } """ import torch.nn as nn class Conv3DSimple(nn.Conv3d): def __init__(self, in_planes, out_planes, midplanes=None, stride=1, padding=1): super().__init__( in_channels=in_planes, out_channels=out_planes, kernel_size=(3, 3, 3), stride=stride, padding=padding, bias=False) @staticmethod def get_downsample_stride(stride): return stride, stride, stride class BasicBlock(nn.Module): expansion = 1 def __init__(self, inplanes, planes, conv_builder, stride=1, downsample=None): midplanes = (inplanes * planes * 3 * 3 * 3) // (inplanes * 3 * 3 + 3 * planes) super(BasicBlock, self).__init__() self.conv1 = nn.Sequential( conv_builder(inplanes, planes, midplanes, stride), nn.BatchNorm3d(planes), nn.ReLU(inplace=True) ) self.conv2 = nn.Sequential( conv_builder(planes, planes, midplanes), nn.BatchNorm3d(planes) ) self.relu = nn.ReLU(inplace=True) self.downsample = downsample self.stride = stride def forward(self, x): residual = x out = self.conv1(x) out = self.conv2(out) if self.downsample is not None: residual = self.downsample(x) out += residual out = self.relu(out) return out class Bottleneck(nn.Module): expansion = 4 def __init__(self, inplanes, planes, conv_builder, stride=1, downsample=None): super(Bottleneck, self).__init__() midplanes = (inplanes * planes * 3 * 3 * 3) // (inplanes * 3 * 3 + 3 * planes) # 1x1x1 self.conv1 = nn.Sequential( nn.Conv3d(inplanes, planes, kernel_size=1, bias=False), nn.BatchNorm3d(planes), nn.ReLU(inplace=True) ) # Second kernel self.conv2 = nn.Sequential( conv_builder(planes, planes, midplanes, stride), nn.BatchNorm3d(planes), nn.ReLU(inplace=True) ) # 1x1x1 self.conv3 = nn.Sequential( nn.Conv3d(planes, planes * self.expansion, kernel_size=1, bias=False), nn.BatchNorm3d(planes * self.expansion) ) self.relu = nn.ReLU(inplace=True) self.downsample = downsample self.stride = stride def forward(self, x): residual = x out = self.conv1(x) out = self.conv2(out) out = self.conv3(out) if self.downsample is not None: residual = self.downsample(x) out += residual out = self.relu(out) return out class BasicStem(nn.Sequential): """The default conv-batchnorm-relu stem """ def __init__(self): super(BasicStem, self).__init__( nn.Conv3d(1, 64, kernel_size=(7, 7, 7), stride=(2, 2, 2), padding=(3, 3, 3), bias=False), nn.BatchNorm3d(64), nn.ReLU(inplace=True)) class VideoResNet(nn.Module): def __init__(self, block, conv_makers, layers, stem, num_classes=2, zero_init_residual=False): """Generic resnet video generator. Args: block (nn.Module): resnet building block conv_makers (list(functions)): generator function for each layer layers (List[int]): number of blocks per layer stem (nn.Module, optional): Resnet stem, if None, defaults to conv-bn-relu. Defaults to None. num_classes (int, optional): Dimension of the final FC layer. Defaults to 400. zero_init_residual (bool, optional): Zero init bottleneck residual BN. Defaults to False. """ super(VideoResNet, self).__init__() self.inplanes = 64 self.stem = stem() self.layer1 = self._make_layer(block, conv_makers[0], 64, layers[0], stride=1) self.layer2 = self._make_layer(block, conv_makers[1], 128, layers[1], stride=2) self.layer3 = self._make_layer(block, conv_makers[2], 192, layers[2], stride=2) self.layer4 = self._make_layer(block, conv_makers[3], 256, layers[3], stride=2) self.avgpool = nn.AdaptiveAvgPool3d((1, 1, 1)) # self.fc = nn.Linear(256 * block.expansion, num_classes) # init weights self._initialize_weights() if zero_init_residual: for m in self.modules(): if isinstance(m, Bottleneck): nn.init.constant_(m.bn3.weight, 0) def forward(self, x): x = self.stem(x) x = self.layer1(x) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) x = self.avgpool(x) # Flatten the layer to fc x = x.flatten(1) # x = self.fc(x) return x def _make_layer(self, block, conv_builder, planes, blocks, stride=1): downsample = None if stride != 1 or self.inplanes != planes * block.expansion: ds_stride = conv_builder.get_downsample_stride(stride) downsample = nn.Sequential( nn.Conv3d(self.inplanes, planes * block.expansion, kernel_size=1, stride=ds_stride, bias=False), nn.BatchNorm3d(planes * block.expansion) ) layers = [] layers.append(block(self.inplanes, planes, conv_builder, stride, downsample)) self.inplanes = planes * block.expansion for i in range(1, blocks): layers.append(block(self.inplanes, planes, conv_builder)) return nn.Sequential(*layers) def _initialize_weights(self): for m in self.modules(): if isinstance(m, nn.Conv3d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.BatchNorm3d): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) elif isinstance(m, nn.Linear): nn.init.normal_(m.weight, 0, 0.01) nn.init.constant_(m.bias, 0) def _video_resnet(arch, pretrained=False, progress=True, **kwargs): model = VideoResNet(**kwargs) return model def r3d_18(pretrained=False, progress=True, **kwargs): """Construct 18 layer Resnet3D model as in https://arxiv.org/abs/1711.11248 Args: pretrained (bool): If True, returns a model pre-trained on Kinetics-400 progress (bool): If True, displays a progress bar of the download to stderr Returns: nn.Module: R3D-18 network """ return _video_resnet('r3d_18', pretrained, progress, block=BasicBlock, conv_makers=[Conv3DSimple] * 4, layers=[2, 2, 2, 2], stem=BasicStem, **kwargs) if __name__ == '__main__': """ ... """ import torch net = r3d_18().to('cuda:1') x = torch.zeros(8, 1, 182, 218, 182).to('cuda:1') print(net(x).shape)