code: # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
  
  import math
  from typing import Optional
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from pytorchvideo.layers.utils import set_attributes
  
  
  class MemoryBank(nn.Module):
      """
      Performs Non-Parametric Instance Discrimination for self supervised learning on
      video. A memory bank is built to keep and update the historical feature embedding
      and use them for contrastive learning.
  
      The original paper is:
      Unsupervised Feature Learning via Non-Parametric Instance Discrimination
      https://arxiv.org/pdf/1805.01978.pdf
  
      More details can be found from the memory bank part in the following paper:
      Momentum Contrast for Unsupervised Visual Representation Learning
      https://arxiv.org/pdf/1911.05722.pdf
      """
  
      def __init__(
          self,
          backbone: nn.Module,
          mlp: Optional[nn.Module] = None,
          neg_size: int = 4096,
          temperature: float = 0.07,
          bank_size: int = 1280000,
          dim: int = 2048,
          mmt: float = 0.999,
      ) -> None:
          """
          Args:
              backbone (nn.Module): backbone used to forward the input.
              mlp (nn.Module): multi-layer perception used in memory bank instance
                  discrimination model.
              neg_size (int): size of negative samples per instance.
              temperature (float): temperature to use for contrastive learning.
              bank_size (int): size of the memory bank, expected to be the same size as
                  the training set.
              dim (int): dimension of the channel.
              mmt (float): momentum to use.
          """
          super().__init__()
          set_attributes(self, locals())
          self._init_mem_bank(bank_size, dim)
  
      def _init_mem_bank(self, bank_size: int, dim: int) -> None:
          """
          Given the memory bank size and the channel dimension, initialize the memory
              bank.
          Args:
              bank_size (int): size of the memory bank, expected to be the same size as
                   the training set.
              dim (int): dimension of the channel.
          """
          stdv = 1.0 / math.sqrt(dim / 3)
          self.register_buffer(
              "memory",
              torch.rand(
                  bank_size,
                  dim,
              )
              .mul_(2 * stdv)
              .add_(-stdv)
              .to(next(self.backbone.parameters()).device),
          )
  
      def forward(self, x: torch.Tensor, x_ind: torch.Tensor) -> torch.Tensor:
          """
          Perform contrastive learning with random sampled negative instance from the
              memory bank. During training, update the memory bank with latest feature
              embedding.
          Args:
              x (torch.tensor): a batch of image with augmentation. The input tensor
                  shape should able to be feed into the backbone.
              x_ind (torch.tensor): the index of the image x from the dataset. Expected
                  shape is B.
          """
          batch_size = x.shape[0]
          x = self.backbone(x)
          if self.mlp is not None:
              x = self.mlp(x)
          # Normalize the output embedding before multiplication.
          x = F.normalize(x, p=2, dim=1)
          # Random sample negative instances from the memory bank.
          idx = torch.randint(0, self.bank_size, size=(batch_size, self.neg_size + 1)).to(
              x.device
          )
          # Fill the first with positive instances.
          idx.select(1, 0).copy_(x_ind.data)
          weight = torch.index_select(self.memory, 0, idx.view(-1)).detach()
          weight = weight.view(batch_size, self.neg_size + 1, self.dim)
          # Multiplication for contrastive learning.
          out = torch.einsum("bkc,bc->bk", weight, x)
          out = torch.div(out, self.temperature)
          gt = torch.zeros((batch_size,), device=x.device, dtype=torch.long)
          loss = torch.nn.functional.cross_entropy(out, gt)
          # Update memory during training.
          if self.training:
              with torch.no_grad():
                  pos = torch.index_select(self.memory, 0, x_ind.view(-1))
                  pos.mul_(self.mmt)
                  pos.add_(torch.mul(x, 1 - self.mmt))
                  norm = pos.pow(2).sum(1, keepdim=True).pow(0.5)
                  updated = pos.div(norm)
                  self.memory.index_copy_(0, x_ind, updated)
          return loss
   # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
  
  import itertools
  import unittest
  
  import numpy as np
  import torch
  from pytorchvideo.layers.convolutions import (
      Conv2plus1d,
      ConvReduce3D,
      create_conv_2plus1d,
  )
  from torch import nn
  
  
  class TestConvReduce3D(unittest.TestCase):
      def setUp(self):
          super().setUp()
          torch.set_rng_state(torch.manual_seed(42).get_state())
  
      def test_create_stack_conv(self):
          """
          Test ConvReduce3D.
          """
          for input_dim, output_dim in itertools.product((2, 4), (4, 8, 16)):
              model = ConvReduce3D(
                  in_channels=input_dim,
                  out_channels=output_dim,
                  kernel_size=((1, 1, 1), (3, 3, 3), (1, 3, 3)),
                  stride=((1, 1, 1), (1, 1, 1), None),
                  padding=((0, 0, 0), (1, 1, 1), (0, 1, 1)),
                  dilation=((2, 2, 2), (1, 1, 1), None),
                  groups=(1, 2, None),
                  bias=(True, False, None),
              )
              model_gt_list = [
                  nn.Conv3d(
                      in_channels=input_dim,
                      out_channels=output_dim,
                      kernel_size=(1, 1, 1),
                      stride=(1, 1, 1),
                      padding=(0, 0, 0),
                      dilation=(2, 2, 2),
                      groups=1,
                      bias=True,
                  ),
                  nn.Conv3d(
                      in_channels=input_dim,
                      out_channels=output_dim,
                      kernel_size=(3, 3, 3),
                      stride=(1, 1, 1),
                      padding=(1, 1, 1),
                      dilation=(1, 1, 1),
                      groups=2,
                      bias=False,
                  ),
                  nn.Conv3d(
                      in_channels=input_dim,
                      out_channels=output_dim,
                      kernel_size=(1, 3, 3),
                      padding=(0, 1, 1),
                  ),
              ]
              model.convs[0].load_state_dict(
                  model_gt_list[0].state_dict(), strict=True
              )  # explicitly use strict mode.
              model.convs[1].load_state_dict(
                  model_gt_list[1].state_dict(), strict=True
              )  # explicitly use strict mode.
              model.convs[2].load_state_dict(
                  model_gt_list[2].state_dict(), strict=True
              )  # explicitly use strict mode.
  
              # Test forwarding.
              for tensor in TestConvReduce3D._get_inputs(input_dim):
                  if tensor.shape[1] != input_dim:
                      with self.assertRaises(RuntimeError):
                          output_tensor = model(tensor)
                      continue
                  else:
                      output_tensor = model(tensor)
                      output_gt = []
                      for ind in range(3):
                          output_gt.append(model_gt_list[ind](tensor))
                      output_tensor_gt = torch.stack(output_gt, dim=0).sum(
                          dim=0, keepdim=False
                      )
  
                  self.assertEqual(
                      output_tensor.shape,
                      output_tensor_gt.shape,
                      "Output shape {} is different from expected shape {}".format(
                          output_tensor.shape, output_tensor_gt.shape
                      ),
                  )
  
      @staticmethod
      def _get_inputs(input_dim: int = 3) -> torch.tensor:
          """
          Provide different tensors as test cases.
  
          Yield:
              (torch.tensor): tensor as test case input.
          """
          # Prepare random tensor as test cases.
          shapes = (
              # Forward succeeded.
              (1, input_dim, 3, 7, 7),
              (1, input_dim, 5, 7, 7),
              (1, input_dim, 7, 7, 7),
              (2, input_dim, 3, 7, 7),
              (4, input_dim, 3, 7, 7),
              (8, input_dim, 3, 7, 7),
              (2, input_dim, 3, 7, 14),
              (2, input_dim, 3, 14, 7),
              (2, input_dim, 3, 14, 14),
              # Forward failed.
              (8, input_dim * 2, 3, 7, 7),
              (8, input_dim * 4, 5, 7, 7),
          )
          for shape in shapes:
              yield torch.rand(shape)
  
  
  class TestConv2plus1d(unittest.TestCase):
      def setUp(self):
          super().setUp()
          torch.set_rng_state(torch.manual_seed(42).get_state())
  
      def test_create_2plus1d_conv(self):
          """
          Test Conv2plus1d.
          """
          for input_dim, output_dim in itertools.product((2, 4), (4, 8, 16)):
              model = Conv2plus1d(
                  conv_t=nn.Conv3d(
                      in_channels=input_dim,
                      out_channels=output_dim,
                      kernel_size=(3, 1, 1),
                      stride=(2, 1, 1),
                      padding=(1, 0, 0),
                      bias=False,
                  ),
                  norm=nn.BatchNorm3d(output_dim),
                  activation=nn.ReLU(),
                  conv_xy=nn.Conv3d(
                      in_channels=output_dim,
                      out_channels=output_dim,
                      kernel_size=(1, 3, 3),
                      stride=(1, 2, 2),
                      padding=(0, 1, 1),
                      bias=False,
                  ),
              )
  
              model_gt = create_conv_2plus1d(
                  in_channels=input_dim,
                  out_channels=output_dim,
                  kernel_size=(3, 3, 3),
                  stride=(2, 2, 2),
                  padding=(1, 1, 1),
                  bias=False,
                  norm=nn.BatchNorm3d,
                  norm_eps=1e-5,
                  norm_momentum=0.1,
                  activation=nn.ReLU,
              )
  
              model.load_state_dict(
                  model_gt.state_dict(), strict=True
              )  # explicitly use strict mode.
  
              # Test forwarding.
              for input_tensor in TestConv2plus1d._get_inputs():
                  with torch.no_grad():
                      if input_tensor.shape[1] != input_dim:
                          with self.assertRaises(RuntimeError):
                              output_tensor = model(input_tensor)
                          continue
                      else:
                          output_tensor = model(input_tensor)
                          output_tensor_gt = model_gt(input_tensor)
                  self.assertEqual(
                      output_tensor.shape,
                      output_tensor_gt.shape,
                      "Output shape {} is different from expected shape {}".format(
                          output_tensor.shape, output_tensor_gt.shape
                      ),
                  )
                  self.assertTrue(
                      np.allclose(output_tensor.numpy(), output_tensor_gt.numpy())
                  )
  
      @staticmethod
      def _get_inputs(input_dim: int = 3) -> torch.tensor:
          """
          Provide different tensors as test cases.
  
          Yield:
              (torch.tensor): tensor as test case input.
          """
          # Prepare random tensor as test cases.
          shapes = (
              # Forward succeeded.
              (1, input_dim, 3, 7, 7),
              (1, input_dim, 5, 7, 7),
              (1, input_dim, 7, 7, 7),
              (2, input_dim, 3, 7, 7),
              (4, input_dim, 3, 7, 7),
              (8, input_dim, 3, 7, 7),
              (2, input_dim, 3, 7, 14),
              (2, input_dim, 3, 14, 7),
              (2, input_dim, 3, 14, 14),
              # Forward failed.
              (8, input_dim * 2, 3, 7, 7),
              (8, input_dim * 4, 5, 7, 7),
          )
          for shape in shapes:
              yield torch.rand(shape)
  
apis: torch.randint torch.nn.functional.normalize torch.rand torch.no_grad torch.div torch.mul torch.nn.functional.cross_entropy torch.zeros torch.einsum
apis: torch.stack torch.nn.BatchNorm3d torch.rand torch.manual_seed torch.no_grad torch.nn.ReLU torch.nn.Conv3d