class_docstr: Contrastive Model, currently mainly focused on memory bank and CSC.
class_docstr_tok: ["Contrastive","Model",",","currently","mainly","focused","on","memory","bank","and","CSC","."]
code: def forward(
      self, clips, index=None, time=None, epoch_exact=None, keys=None
  ):
      if epoch_exact is not None and self.momentum_annealing:
          self.momentum_anneal_cosine(epoch_exact)
  
      if self.type == "mem":
          batch_size = clips[0].size(0)
          q = self.backbone(clips)
          if index is None:
              return q
          q = self.l2_norm(q)
  
          if not self.training:
              assert self.knn_mem.duration == 1
              return self.eval_knn(q)
          time *= self.duration - 1
          clip_ind = torch.randint(
              0,
              self.length,
              size=(
                  batch_size,
                  self.k + 1,
              ),
          ).cuda()
          clip_ind.select(1, 0).copy_(index.data)
  
          if self.mem_type == "2d":
              if self.interp:
                  time_ind = (
                      torch.empty(batch_size, self.k + 1)
                      .uniform_(0, self.duration - 1)
                      .cuda()
                  )
              else:
                  time_ind = torch.randint(
                      0,
                      self.duration - 1,
                      size=(
                          batch_size,
                          self.k + 1,
                      ),
                  ).cuda()
          else:
              time_ind = torch.zeros(
                  size=(batch_size, self.k + 1), dtype=int
              ).cuda()
  
          if self.examplar_type == "clip":
              # Diff clip from same video are negative.
              time_ind.select(1, 0).copy_(time.data)
          elif self.examplar_type == "video":
              # Diff clip from same video are positive.
              pass
          else:
              raise NotImplementedError(
                  "unsupported examplar_type {}".format(
                      self.examplar_type
                  )
              )
          k = self.memory.get(clip_ind, time_ind, self.interp)
          # q: N x C, k: N x K x C
          prod = torch.einsum("nc,nkc->nk", q, k)
          prod = torch.div(prod, self.T)
  
          loss = self.nce_loss_fun(prod)
  
          self.memory.update(
              q,
              momentum=self.mmt,
              ind=index,
              time=time,
              interp=self.interp,
          )
          self.knn_mem_update(q, index)
          return prod, 0.0, True
      elif self.type == "moco":
          if isinstance(clips[0], list):
              n_clips = len(clips)
              ind_clips = np.arange(
                  n_clips
              )  # clips come ordered temporally from decoder
  
              clip_q = clips[ind_clips[0]]
              clips_k = [clips[i] for i in ind_clips[1:]]
              # rearange time
              time_q = time[:, ind_clips[0], :]
              time_k = (
                  time[:, ind_clips[1:], :]
                  if keys is None
                  else time[:, ind_clips[0] + 1 :, :]
              )
          else:
              clip_q = clips
  
          feat_q = self.backbone(clip_q)
          extra_projs = []
          if isinstance(feat_q, list):
              extra_projs = feat_q[1:]
              feat_q = feat_q[0]
              extra_projs = [self.l2_norm(feat) for feat in extra_projs]
  
          if index is None:
              return feat_q
          q = self.l2_norm(feat_q)
          q_knn = q
  
          if not self.training:
              return self.eval_knn(q_knn)
  
          if keys is None:
              keys = self.compute_key_feat(
                  clips_k, compute_predictor_keys=False
              )
              auto_enqueue_keys = True
          else:
              auto_enqueue_keys = False
  
          # score computation
          queue_neg = torch.einsum(
              "nc,kc->nk", [q, self.queue_x.clone().detach()]
          )
  
          for k, key in enumerate(keys):
              out_pos = torch.einsum("nc,nc->n", [q, key]).unsqueeze(-1)
              lgt_k = torch.cat([out_pos, queue_neg], dim=1)
              if k == 0:
                  logits = lgt_k
              else:
                  logits = torch.cat([logits, lgt_k], dim=0)
  
          logits = torch.div(logits, self.T)
  
          loss = self.nce_loss_fun(logits)
          # update queue
          if self.training and auto_enqueue_keys:
              self._dequeue_and_enqueue(keys)
  
          self.knn_mem_update(q_knn, index)
          return logits, loss
  
      elif self.type == "byol":
          clips_key = [None] * len(clips)
          for i, clip in enumerate(clips):
              p = []
              for path in clip:
                  p.append(path)
              clips_key[i] = p
          batch_clips = False
          if isinstance(clips[0], list):
              n_clips = len(clips)
              ind_clips = np.arange(
                  n_clips
              )  # clips come ordered temporally from decoder
              if batch_clips and n_clips > 1:
                  clips_batched = self.batch_clips(clips)
                  clips_key = [clips_batched]
                  clip_q = clips_batched
              else:
                  clip_q = clips[0]
          else:
              clip_q = clips
  
          feat_q = self.backbone(clip_q)
          predictors = []
          if isinstance(feat_q, list):
              predictors = feat_q[1:]
              feat_q = feat_q[0]
              predictors = [self.l2_norm(feat) for feat in predictors]
          else:
              raise NotImplementedError("BYOL: predictor is missing")
          assert len(predictors) == 1
          if index is None:
              return feat_q
          q = self.l2_norm(feat_q)
  
          q_knn = q  # projector
  
          if not self.training:
              return self.eval_knn(q_knn)
  
          ind_clips = np.arange(
              n_clips
          )  # clips come ordered temporally from decoder
  
          # rest down is for training
          if keys is None:
              keys = self.compute_key_feat(
                  clips_key, compute_predictor_keys=False
              )
  
          if self.cfg.CONTRASTIVE.SEQUENTIAL:
              loss_reg = self.sim_loss(predictors[0], keys[0])
              for i in range(1, len(keys)):
                  loss_reg += self.sim_loss(predictors[0], keys[i])
              loss_reg /= len(keys)
          else:
              if batch_clips:
                  bs = predictors[0].shape[0] // 2
                  loss_reg = self.sim_loss(
                      predictors[0][:bs, :], keys[0][bs:, :]
                  ) + self.sim_loss(
                      predictors[0][bs:, :], keys[0][:bs, :]
                  )
                  q_knn = q_knn[:bs, :]
                  del clips_batched[0]
              else:
                  loss_q1 = self.sim_loss(predictors[0], keys[1])
                  assert len(clips) == 2
                  clip_q2 = clips[1]
                  feat_q2 = self.backbone(clip_q2)
                  predictors2 = feat_q2[1:]
                  # feat_q2 = feat_q2[0] # not used
                  predictors2 = [
                      self.l2_norm(feat) for feat in predictors2
                  ]
                  assert len(predictors2) == 1
  
                  loss_q2 = self.sim_loss(predictors2[0], keys[0])
                  loss_reg = loss_q1 + loss_q2
  
          # loss_pos = self.sim_loss(q1_proj, q2_proj)
          dummy_logits = torch.cat(
              (
                  9999.0
                  * torch.ones(
                      (len(index), 1), dtype=torch.float
                  ).cuda(),
                  torch.zeros(
                      (len(index), self.k), dtype=torch.float
                  ).cuda(),
              ),
              dim=1,
          )
  
          self.knn_mem_update(q_knn, index)
  
          return dummy_logits, loss_reg
  
      elif self.type == "swav":
          if not isinstance(clips[0], list):
              if self.swav_use_public_code:
                  proj_1, _ = self.run_swav_orig_encoder_q(clips)
              else:
                  proj_1, _ = self.run_swav_encoder_q(clips)
              if index is None:
                  return proj_1
              if not self.training:
                  return self.eval_knn(proj_1)
          n_clips = len(clips)
          ind_clips = np.arange(
              n_clips
          )  # clips come ordered temporally from decoder
          clip_q = clips[0]
  
          if self.swav_use_public_code:
              # uses official code of SwAV from
              # https://github.com/facebookresearch/swav/blob/master/main_swav.py
              with torch.no_grad():
                  m = getattr(self, "module", self)
                  w = m.swav_prototypes.weight.data.clone()
                  w = nn.functional.normalize(w, dim=1, p=2)
                  m.swav_prototypes.weight.copy_(w)
  
              bs = clips[0][0].size(0)
              output, embedding = [], []
  
              for i, clip_q in enumerate(clips):
                  x = self.run_swav_orig_encoder_q(clip_q)
                  embedding.append(x[0])
                  output.append(x[1])
              q_knn = embedding[0]
              embedding = torch.cat(embedding, dim=0)
              output = torch.cat(output, dim=0)
  
              loss_swav = 0
              swav_extra_crops = n_clips - 2
              self.swav_crops_for_assign = np.arange(
                  n_clips - swav_extra_crops
              )
              for i, crop_id in enumerate(self.swav_crops_for_assign):
                  with torch.no_grad():
                      out = output[bs * crop_id : bs * (crop_id + 1)]
                      if (
                          self.cfg.CONTRASTIVE.SWAV_QEUE_LEN > 0
                          and epoch_exact >= 15.0
                      ):
                          if self.swav_use_the_queue or not torch.all(
                              self.queue_swav[i, -1, :] == 0
                          ):
                              self.swav_use_the_queue = True
                              out = torch.cat(
                                  (
                                      torch.mm(
                                          self.queue_swav[i],
                                          m.swav_prototypes.weight.t(),
                                      ),
                                      out,
                                  )
                              )
                          self.queue_swav[i, bs:] = self.queue_swav[
                              i, :-bs
                          ].clone()
                          self.queue_swav[i, :bs] = embedding[
                              crop_id * bs : (crop_id + 1) * bs
                          ]
                      q = out / self.swav_eps_sinkhorn
                      q = torch.exp(q).t()
                      q = (
                          self.distributed_sinkhorn(q, 3)[-bs:]
                          if self.cfg.NUM_SHARDS > 1
                          else self.sinkhorn(q.t(), 3)[-bs:]
                      )
                  subloss = 0
                  for v in np.delete(np.arange(n_clips), crop_id):
                      p = self.softmax(
                          output[bs * v : bs * (v + 1)] / self.T
                      )
                      subloss -= torch.mean(
                          torch.sum(q * torch.log(p), dim=1)
                      )
                  loss_swav += subloss / (n_clips - 1)
              loss_swav /= len(self.swav_crops_for_assign)
          else:
              proj_1, out_1 = self.run_swav_encoder_q(clip_q)
              q_knn = proj_1
              if not self.training:
                  return self.eval_knn(q_knn)
              proj_2, out_2 = self.run_swav_encoder_q(clips[1])
              bs = proj_1.shape[0]
              if self.cfg.CONTRASTIVE.SWAV_QEUE_LEN > 0:
                  if epoch_exact >= 15.0 and not torch.all(
                      self.queue_swav[0, -1, :] == 0
                  ):
                      swav_prototypes = F.normalize(
                          self.swav_prototypes, dim=0, p=2
                      ).detach()
                      out_1 = torch.cat(
                          (
                              torch.mm(
                                  self.queue_swav[0].detach(),
                                  swav_prototypes,
                              ),
                              out_1,
                          )
                      )
                      out_2 = torch.cat(
                          (
                              torch.mm(
                                  self.queue_swav[1].detach(),
                                  swav_prototypes,
                              ),
                              out_2,
                          )
                      )
                  # fill the queue
                  self.queue_swav[0, bs:] = self.queue_swav[
                      0, :-bs
                  ].clone()
                  self.queue_swav[0, :bs] = proj_1.detach()
                  self.queue_swav[1, bs:] = self.queue_swav[
                      1, :-bs
                  ].clone()
                  self.queue_swav[1, :bs] = proj_2.detach()
  
              with torch.no_grad():
                  code_1 = self.get_code(out_1)
                  code_2 = self.get_code(out_2)
              loss12 = self.KLDivLoss(
                  out_1[-bs:], code_2[-bs:].detach()
              )
              loss21 = self.KLDivLoss(
                  out_2[-bs:], code_1[-bs:].detach()
              )
              loss_swav = loss12 + loss21
          self.knn_mem_update(q_knn, index)
          dummy_logits = torch.cat(
              (
                  9999.0
                  * torch.ones(
                      (len(index), 1), dtype=torch.float
                  ).cuda(),
                  torch.zeros(
                      (len(index), self.k), dtype=torch.float
                  ).cuda(),
              ),
              dim=1,
          )
          return dummy_logits, loss_swav
  
      elif self.type == "simclr":
          if isinstance(clips[0], list):
              n_clips = len(clips)
              clip_q = clips[0]
          else:
              clip_q = clips
          feat_q = self.backbone(clip_q)
          q = self.l2_norm(feat_q)
          if index is None:
              return q
          q_knn = q
          if not self.training:
              return self.eval_knn(q_knn)
          q2 = self.backbone(clips[1])
          q2 = self.l2_norm(q2)
          distributed_loss = False
          if distributed_loss and self.cfg.NUM_GPUS > 1:
              out = torch.cat([q, q2], dim=0)
              if self.cfg.CONTRASTIVE.SIMCLR_DIST_ON:
                  out_all = du.cat_all_gather(out)
              else:
                  out_all = out
              similarity = torch.exp(
                  torch.mm(out, out_all.t()) / self.T
              )
              Z, loss = 0.0, 0.0
              for loss_id in range(len(self.pos_mask)):
                  pos = torch.sum(
                      similarity * self.pos_mask[loss_id], 1
                  )
                  neg = torch.sum(similarity * self.neg_mask, 1)
                  idx = (
                      1 - torch.sum(self.pos_mask[loss_id], 1) > 0
                  ).detach()
                  term_prob = pos / (pos + neg)
                  term_prob[idx] = 1.0
                  term_loss = torch.log(term_prob)
                  Z += torch.sum(~idx).detach()
                  loss -= torch.sum(term_loss)
              loss /= Z
          else:
              cat_across_gpus = True
              if cat_across_gpus and self.cfg.NUM_GPUS > 1:
                  q = du.AllGatherWithGradient.apply(q)
                  q2 = du.AllGatherWithGradient.apply(q2)
              out = torch.cat([q, q2], dim=0)
              # [2*B, 2*B]
              sim_matrix = torch.exp(
                  torch.mm(out, out.t().contiguous()) / self.T
              )
              # SANITY:
              mask = (
                  torch.ones_like(sim_matrix)
                  - torch.eye(out.shape[0], device=sim_matrix.device)
              ).bool()
              # [2*B, 2*B-1]
              sim_matrix = sim_matrix.masked_select(mask).view(
                  out.shape[0], -1
              )
              # compute loss
              pos_sim = torch.exp(torch.sum(q * q2, dim=-1) / self.T)
              # [2*B]
              pos_sim = torch.cat([pos_sim, pos_sim], dim=0)
              loss = (
                  -torch.log(pos_sim / sim_matrix.sum(dim=-1))
              ).mean()
          self.knn_mem_update(q_knn, index)
          dummy_logits = torch.cat(
              (
                  9999.0
                  * torch.ones(
                      (len(index), 1), dtype=torch.float
                  ).cuda(),
                  torch.zeros(
                      (len(index), self.k), dtype=torch.float
                  ).cuda(),
              ),
              dim=1,
          )
          return dummy_logits, loss
      else:
          raise NotImplementedError()
  
imports: import math
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import slowfast.models.losses
  import slowfast.utils.distributed
  import slowfast.utils.logging
  from slowfast.models.video_model_builder import X3D, MViT, ResNet, SlowFast
  from .build import MODEL_REGISTRY
label_desc: Use image node_id 13 for calling the ContrastiveModel obj's underlying member method code with example usage: obj.forward(clips, index, time, epoch_exact, keys) and returns: prod,  int,  True, q, self, logits,  loss, feat_q, self, dummy_logits,  loss_reg, feat_q, self, dummy_logits,  loss_swav, dummy_logits,  loss, proj_1, self, self, q, self