text: #!/usr/bin/env python3
  # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
  
  """Video models."""
  
  import torch.nn as nn
  
  from slowfast.models.nonlocal_helper import Nonlocal
  
  
  import math
  
  import torch.nn as nn
  from torch.nn.modules.utils import _triple
  
  
  
  
  def get_trans_func(name):
      """
      Retrieves the transformation module by name.
      """
      trans_funcs = {
          "bottleneck_transform": BottleneckTransform,
          "basic_transform": BasicTransform,
      }
      assert (
          name in trans_funcs.keys()
      ), "Transformation function '{}' not supported".format(name)
      return trans_funcs[name]
  
  
  # This class is based on https://github.com/irhum/R2Plus1D-PyTorch
  # based on the work (2+1)R https://arxiv.org/pdf/1711.11248.pdf
  class SpatioTemporalConv(nn.Module):
      r"""Applies a factored 3D convolution over an input signal composed of several input
      planes with distinct spatial and time axes, by performing a 2D convolution over the
      spatial axes to an intermediate subspace, followed by a 1D convolution over the time
      axis to produce the final output.
      Args:
          in_channels (int): Number of channels in the input tensor
          out_channels (int): Number of channels produced by the convolution
          kernel_size (int or tuple): Size of the convolving kernel
          stride (int or tuple, optional): Stride of the convolution. Default: 1
          padding (int or tuple, optional): Zero-padding added to the sides of the input during their respective convolutions. Default: 0
          bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``
      """
  
      def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, bias=True, **kwargs):
          super(SpatioTemporalConv, self).__init__()
  
          # if ints are entered, convert them to iterables, 1 -> [1, 1, 1]
          kernel_size = _triple(kernel_size)
          stride = _triple(stride)
          padding = _triple(padding)
  
          # decomposing the parameters into spatial and temporal components by
          # masking out the values with the defaults on the axis that
          # won't be convolved over. This is necessary to avoid unintentional
          # behavior such as padding being added twice
          spatial_kernel_size =  [1, kernel_size[1], kernel_size[2]]
          spatial_stride =  [1, stride[1], stride[2]]
          spatial_padding =  [0, padding[1], padding[2]]
  
          temporal_kernel_size = [kernel_size[0], 1, 1]
          temporal_stride =  [stride[0], 1, 1]
          temporal_padding =  [padding[0], 0, 0]
  
          # compute the number of intermediary channels (M) using formula
          # from the paper section 3.5
          intermed_channels = int(math.floor((kernel_size[0] * kernel_size[1] * kernel_size[2] * in_channels * out_channels)/ \
                              (kernel_size[1]* kernel_size[2] * in_channels + kernel_size[0] * out_channels)))
  
          # the spatial conv is effectively a 2D conv due to the
          # spatial_kernel_size, followed by batch_norm and ReLU
          self.spatial_conv = nn.Conv3d(in_channels, intermed_channels, spatial_kernel_size,
                                      stride=spatial_stride, padding=spatial_padding, bias=bias, **kwargs)
          self.bn = nn.BatchNorm3d(intermed_channels)
          self.relu = nn.ReLU()
  
          # the temporal conv is effectively a 1D conv, but has batch norm
          # and ReLU added inside the model constructor, not here. This is an
          # intentional design choice, to allow this module to externally act
          # identical to a standard Conv3D, so it can be reused easily in any
          # other codebase
          self.temporal_conv = nn.Conv3d(intermed_channels, out_channels, temporal_kernel_size,
                                         stride=temporal_stride, padding=temporal_padding, bias=bias, **kwargs)
  
      def forward(self, x):
          x = self.relu(self.bn(self.spatial_conv(x)))
          x = self.temporal_conv(x)
          return x
  
  
  class BasicTransform(nn.Module):
      """
      Basic transformation: Tx3x3, 1x3x3, where T is the size of temporal kernel.
      """
  
      def __init__(
          self,
          dim_in,
          dim_out,
          temp_kernel_size,
          stride,
          dim_inner=None,
          num_groups=1,
          stride_1x1=None,
          inplace_relu=True,
          eps=1e-5,
          bn_mmt=0.1,
          norm_module=nn.BatchNorm3d,
          type_conv='3d'
      ):
          """
          Args:
              dim_in (int): the channel dimensions of the input.
              dim_out (int): the channel dimension of the output.
              temp_kernel_size (int): the temporal kernel sizes of the middle
                  convolution in the bottleneck.
              stride (int): the stride of the bottleneck.
              dim_inner (None): the inner dimension would not be used in
                  BasicTransform.
              num_groups (int): number of groups for the convolution. Number of
                  group is always 1 for BasicTransform.
              stride_1x1 (None): stride_1x1 will not be used in BasicTransform.
              inplace_relu (bool): if True, calculate the relu on the original
                  input without allocating new memory.
              eps (float): epsilon for batch norm.
              bn_mmt (float): momentum for batch norm. Noted that BN momentum in
                  PyTorch = 1 - BN momentum in Caffe2.
              norm_module (nn.Module): nn.Module for the normalization layer. The
                  default is nn.BatchNorm3d.
          """
          super(BasicTransform, self).__init__()
          self.temp_kernel_size = temp_kernel_size
          self._inplace_relu = inplace_relu
          self._eps = eps
          self._bn_mmt = bn_mmt
  
          self._type_conv = nn.Conv3d if type_conv == '3d' else SpatioTemporalConv
          self._construct(dim_in, dim_out, stride, norm_module)
  
      def _construct(self, dim_in, dim_out, stride, norm_module):
          # Tx3x3, BN, ReLU.
          self.a = self._type_conv(#SpatioTemporalConv(#nn.Conv3d(
              dim_in,
              dim_out,
              kernel_size=[self.temp_kernel_size, 3, 3],
              stride=[1, stride, stride],
              padding=[int(self.temp_kernel_size // 2), 1, 1],
              bias=False,
          )
          self.a_bn = norm_module(
              num_features=dim_out, eps=self._eps, momentum=self._bn_mmt
          )
          self.a_relu = nn.ReLU(inplace=self._inplace_relu)
          # 1x3x3, BN.
          self.b = self._type_conv(#SpatioTemporalConv(#nn.Conv3d(
              dim_out,
              dim_out,
              kernel_size=[1, 3, 3],
              stride=[1, 1, 1],
              padding=[0, 1, 1],
              bias=False,
          )
          self.b_bn = norm_module(
              num_features=dim_out, eps=self._eps, momentum=self._bn_mmt
          )
  
          self.b_bn.transform_final_bn = True
  
      def forward(self, x):
          x = self.a(x)
          x = self.a_bn(x)
          x = self.a_relu(x)
  
          x = self.b(x)
          x = self.b_bn(x)
          return x
  
  
  class BottleneckTransform(nn.Module):
      """
      Bottleneck transformation: Tx1x1, 1x3x3, 1x1x1, where T is the size of
          temporal kernel.
      """
  
      def __init__(
          self,
          dim_in,
          dim_out,
          temp_kernel_size,
          stride,
          dim_inner,
          num_groups,
          stride_1x1=False,
          inplace_relu=True,
          eps=1e-5,
          bn_mmt=0.1,
          dilation=1,
          norm_module=nn.BatchNorm3d,
          type_conv='3d'
      ):
          """
          Args:
              dim_in (int): the channel dimensions of the input.
              dim_out (int): the channel dimension of the output.
              temp_kernel_size (int): the temporal kernel sizes of the middle
                  convolution in the bottleneck.
              stride (int): the stride of the bottleneck.
              dim_inner (int): the inner dimension of the block.
              num_groups (int): number of groups for the convolution. num_groups=1
                  is for standard ResNet like networks, and num_groups>1 is for
                  ResNeXt like networks.
              stride_1x1 (bool): if True, apply stride to 1x1 conv, otherwise
                  apply stride to the 3x3 conv.
              inplace_relu (bool): if True, calculate the relu on the original
                  input without allocating new memory.
              eps (float): epsilon for batch norm.
              bn_mmt (float): momentum for batch norm. Noted that BN momentum in
                  PyTorch = 1 - BN momentum in Caffe2.
              dilation (int): size of dilation.
              norm_module (nn.Module): nn.Module for the normalization layer. The
                  default is nn.BatchNorm3d.
          """
          super(BottleneckTransform, self).__init__()
          self.temp_kernel_size = temp_kernel_size
          self._inplace_relu = inplace_relu
          self._eps = eps
          self._bn_mmt = bn_mmt
          self._stride_1x1 = stride_1x1
          self._type_conv = nn.Conv3d if type_conv == '3d' else SpatioTemporalConv
          self._construct(
              dim_in,
              dim_out,
              stride,
              dim_inner,
              num_groups,
              dilation,
              norm_module,
          )
  
      def _construct(
          self,
          dim_in,
          dim_out,
          stride,
          dim_inner,
          num_groups,
          dilation,
          norm_module,
      ):
          (str1x1, str3x3) = (stride, 1) if self._stride_1x1 else (1, stride)
  
          # Tx1x1, BN, ReLU.
          self.a = self._type_conv(#SpatioTemporalConv(#nn.Conv3d(
              dim_in,
              dim_inner,
              kernel_size=[self.temp_kernel_size, 1, 1],
              stride=[1, str1x1, str1x1],
              padding=[int(self.temp_kernel_size // 2), 0, 0],
              bias=False,
          )
          self.a_bn = norm_module(
              num_features=dim_inner, eps=self._eps, momentum=self._bn_mmt
          )
          self.a_relu = nn.ReLU(inplace=self._inplace_relu)
  
          # 1x3x3, BN, ReLU.
          self.b = self._type_conv(#SpatioTemporalConv(#nn.Conv3d(
              dim_inner,
              dim_inner,
              [1, 3, 3],
              stride=[1, str3x3, str3x3],
              padding=[0, dilation, dilation],
              groups=num_groups,
              bias=False,
              dilation=[1, dilation, dilation],
          )
          self.b_bn = norm_module(
              num_features=dim_inner, eps=self._eps, momentum=self._bn_mmt
          )
          self.b_relu = nn.ReLU(inplace=self._inplace_relu)
  
          # 1x1x1, BN.
          self.c = self._type_conv(#SpatioTemporalConv(#nn.Conv3d(
              dim_inner,
              dim_out,
              kernel_size=[1, 1, 1],
              stride=[1, 1, 1],
              padding=[0, 0, 0],
              bias=False,
          )
          self.c_bn = norm_module(
              num_features=dim_out, eps=self._eps, momentum=self._bn_mmt
          )
          self.c_bn.transform_final_bn = True
  
      def forward(self, x):
          # Explicitly forward every layer.
          # Branch2a.
          x = self.a(x)
          x = self.a_bn(x)
          x = self.a_relu(x)
  
          # Branch2b.
          x = self.b(x)
          x = self.b_bn(x)
          x = self.b_relu(x)
  
          # Branch2c
          x = self.c(x)
          x = self.c_bn(x)
          return x
  
  
  class ResBlock(nn.Module):
      """
      Residual block.
      """
  
      def __init__(
          self,
          dim_in,
          dim_out,
          temp_kernel_size,
          stride,
          trans_func,
          dim_inner,
          num_groups=1,
          stride_1x1=False,
          inplace_relu=True,
          eps=1e-5,
          bn_mmt=0.1,
          dilation=1,
          norm_module=nn.BatchNorm3d,
          type_conv='3d'
      ):
          """
          ResBlock class constructs redisual blocks. More details can be found in:
              Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.
              "Deep residual learning for image recognition."
              https://arxiv.org/abs/1512.03385
          Args:
              dim_in (int): the channel dimensions of the input.
              dim_out (int): the channel dimension of the output.
              temp_kernel_size (int): the temporal kernel sizes of the middle
                  convolution in the bottleneck.
              stride (int): the stride of the bottleneck.
              trans_func (string): transform function to be used to construct the
                  bottleneck.
              dim_inner (int): the inner dimension of the block.
              num_groups (int): number of groups for the convolution. num_groups=1
                  is for standard ResNet like networks, and num_groups>1 is for
                  ResNeXt like networks.
              stride_1x1 (bool): if True, apply stride to 1x1 conv, otherwise
                  apply stride to the 3x3 conv.
              inplace_relu (bool): calculate the relu on the original input
                  without allocating new memory.
              eps (float): epsilon for batch norm.
              bn_mmt (float): momentum for batch norm. Noted that BN momentum in
                  PyTorch = 1 - BN momentum in Caffe2.
              dilation (int): size of dilation.
              norm_module (nn.Module): nn.Module for the normalization layer. The
                  default is nn.BatchNorm3d.
          """
          super(ResBlock, self).__init__()
          self._inplace_relu = inplace_relu
          self._eps = eps
          self._bn_mmt = bn_mmt
          self._type_conv = nn.Conv3d if type_conv == '3d' else SpatioTemporalConv
          self._construct(
              dim_in,
              dim_out,
              temp_kernel_size,
              stride,
              trans_func,
              dim_inner,
              num_groups,
              stride_1x1,
              inplace_relu,
              dilation,
              norm_module,
          )
  
      def _construct(
          self,
          dim_in,
          dim_out,
          temp_kernel_size,
          stride,
          trans_func,
          dim_inner,
          num_groups,
          stride_1x1,
          inplace_relu,
          dilation,
          norm_module,
      ):
          # Use skip connection with projection if dim or res change.
          if (dim_in != dim_out) or (stride != 1):
              self.branch1 = self._type_conv(#SpatioTemporalConv(#nn.Conv3d(
                  dim_in,
                  dim_out,
                  kernel_size=1,
                  stride=[1, stride, stride],
                  padding=0,
                  bias=False,
                  dilation=1,
              )
              self.branch1_bn = norm_module(
                  num_features=dim_out, eps=self._eps, momentum=self._bn_mmt
              )
          self.branch2 = trans_func(
              dim_in,
              dim_out,
              temp_kernel_size,
              stride,
              dim_inner,
              num_groups,
              stride_1x1=stride_1x1,
              inplace_relu=inplace_relu,
              dilation=dilation,
              norm_module=norm_module,
          )
          self.relu = nn.ReLU(self._inplace_relu)
  
      def forward(self, x):
          if hasattr(self, "branch1"):
              x = self.branch1_bn(self.branch1(x)) + self.branch2(x)
          else:
              x = x + self.branch2(x)
          x = self.relu(x)
          return x
  
  
  class ResStage(nn.Module):
      """
      Stage of 3D ResNet. It expects to have one or more tensors as input for
          single pathway (C2D, I3D, Slow), and multi-pathway (SlowFast) cases.
          More details can be found here:
  
          Christoph Feichtenhofer, Haoqi Fan, Jitendra Malik, and Kaiming He.
          "SlowFast networks for video recognition."
          https://arxiv.org/pdf/1812.03982.pdf
      """
  
      def __init__(
          self,
          dim_in,
          dim_out,
          stride,
          temp_kernel_sizes,
          num_blocks,
          dim_inner,
          num_groups,
          num_block_temp_kernel,
          nonlocal_inds,
          nonlocal_group,
          nonlocal_pool,
          dilation,
          instantiation="softmax",
          trans_func_name="bottleneck_transform",
          stride_1x1=False,
          inplace_relu=True,
          norm_module=nn.BatchNorm3d,
          type_conv='3d'
      ):
          """
          The `__init__` method of any subclass should also contain these arguments.
          ResStage builds p streams, where p can be greater or equal to one.
          Args:
              dim_in (list): list of p the channel dimensions of the input.
                  Different channel dimensions control the input dimension of
                  different pathways.
              dim_out (list): list of p the channel dimensions of the output.
                  Different channel dimensions control the input dimension of
                  different pathways.
              temp_kernel_sizes (list): list of the p temporal kernel sizes of the
                  convolution in the bottleneck. Different temp_kernel_sizes
                  control different pathway.
              stride (list): list of the p strides of the bottleneck. Different
                  stride control different pathway.
              num_blocks (list): list of p numbers of blocks for each of the
                  pathway.
              dim_inner (list): list of the p inner channel dimensions of the
                  input. Different channel dimensions control the input dimension
                  of different pathways.
              num_groups (list): list of number of p groups for the convolution.
                  num_groups=1 is for standard ResNet like networks, and
                  num_groups>1 is for ResNeXt like networks.
              num_block_temp_kernel (list): extent the temp_kernel_sizes to
                  num_block_temp_kernel blocks, then fill temporal kernel size
                  of 1 for the rest of the layers.
              nonlocal_inds (list): If the tuple is empty, no nonlocal layer will
                  be added. If the tuple is not empty, add nonlocal layers after
                  the index-th block.
              dilation (list): size of dilation for each pathway.
              nonlocal_group (list): list of number of p nonlocal groups. Each
                  number controls how to fold temporal dimension to batch
                  dimension before applying nonlocal transformation.
                  https://github.com/facebookresearch/video-nonlocal-net.
              instantiation (string): different instantiation for nonlocal layer.
                  Supports two different instantiation method:
                      "dot_product": normalizing correlation matrix with L2.
                      "softmax": normalizing correlation matrix with Softmax.
              trans_func_name (string): name of the the transformation function apply
                  on the network.
              norm_module (nn.Module): nn.Module for the normalization layer. The
                  default is nn.BatchNorm3d.
          """
          super(ResStage, self).__init__()
          assert all(
              (
                  num_block_temp_kernel[i] <= num_blocks[i]
                  for i in range(len(temp_kernel_sizes))
              )
          )
          self.num_blocks = num_blocks
          self.nonlocal_group = nonlocal_group
          self.temp_kernel_sizes = [
              (temp_kernel_sizes[i] * num_blocks[i])[: num_block_temp_kernel[i]]
              + [1] * (num_blocks[i] - num_block_temp_kernel[i])
              for i in range(len(temp_kernel_sizes))
          ]
          assert (
              len(
                  {
                      len(dim_in),
                      len(dim_out),
                      len(temp_kernel_sizes),
                      len(stride),
                      len(num_blocks),
                      len(dim_inner),
                      len(num_groups),
                      len(num_block_temp_kernel),
                      len(nonlocal_inds),
                      len(nonlocal_group),
                  }
              )
              == 1
          )
          self.num_pathways = len(self.num_blocks)
          self._construct(
              dim_in,
              dim_out,
              stride,
              dim_inner,
              num_groups,
              trans_func_name,
              stride_1x1,
              inplace_relu,
              nonlocal_inds,
              nonlocal_pool,
              instantiation,
              dilation,
              norm_module,
              type_conv,
          )
  
      def _construct(
          self,
          dim_in,
          dim_out,
          stride,
          dim_inner,
          num_groups,
          trans_func_name,
          stride_1x1,
          inplace_relu,
          nonlocal_inds,
          nonlocal_pool,
          instantiation,
          dilation,
          norm_module,
          type_conv
      ):
          for pathway in range(self.num_pathways):
              for i in range(self.num_blocks[pathway]):
                  # Retrieve the transformation function.
                  trans_func = get_trans_func(trans_func_name)
                  # Construct the block.
                  res_block = ResBlock(
                      dim_in[pathway] if i == 0 else dim_out[pathway],
                      dim_out[pathway],
                      self.temp_kernel_sizes[pathway][i],
                      stride[pathway] if i == 0 else 1,
                      trans_func,
                      dim_inner[pathway],
                      num_groups[pathway],
                      stride_1x1=stride_1x1,
                      inplace_relu=inplace_relu,
                      dilation=dilation[pathway],
                      norm_module=norm_module,
                      type_conv=type_conv,
                  )
                  self.add_module("pathway{}_res{}".format(pathway, i), res_block)
                  if i in nonlocal_inds[pathway]:
                      nln = Nonlocal(
                          dim_out[pathway],
                          dim_out[pathway] // 2,
                          nonlocal_pool[pathway],
                          instantiation=instantiation,
                          norm_module=norm_module,
                      )
                      self.add_module(
                          "pathway{}_nonlocal{}".format(pathway, i), nln
                      )
  
      def forward(self, inputs):
          output = []
          for pathway in range(self.num_pathways):
              x = inputs[pathway]
              for i in range(self.num_blocks[pathway]):
                  m = getattr(self, "pathway{}_res{}".format(pathway, i))
                  x = m(x)
                  if hasattr(self, "pathway{}_nonlocal{}".format(pathway, i)):
                      nln = getattr(
                          self, "pathway{}_nonlocal{}".format(pathway, i)
                      )
                      b, c, t, h, w = x.shape
                      if self.nonlocal_group[pathway] > 1:
                          # Fold temporal dimension into batch dimension.
                          x = x.permute(0, 2, 1, 3, 4)
                          x = x.reshape(
                              b * self.nonlocal_group[pathway],
                              t // self.nonlocal_group[pathway],
                              c,
                              h,
                              w,
                          )
                          x = x.permute(0, 2, 1, 3, 4)
                      x = nln(x)
                      if self.nonlocal_group[pathway] > 1:
                          # Fold back to temporal dimension.
                          x = x.permute(0, 2, 1, 3, 4)
                          x = x.reshape(b, t, c, h, w)
                          x = x.permute(0, 2, 1, 3, 4)
              output.append(x)
  
          return output
  