class_docstr: Helper to add to hooks to create and update a momentum encoder in the loss
class_docstr_tok: ["Helper","to","add","to","hooks","to","create","and","update","a","momentum","encoder","in","the","loss"]
docstr: Create the teacher: it is an exponential moving average of the student.
docstr_tok: ["Create","the","teacher",":","it","is","an","exponential","moving","average","of","the","student","."]
code: def build_momentum_network(
      task: tasks.ClassyTask,
      config: AttrDict,
  ) -> None:
      """
      Create the teacher: it is an exponential moving average of the student.
      """
      logging.info("Building momentum encoder...")
  
      # Same architecture but do not apply stochastic depth
      is_fsdp_model = "fsdp" in config.MODEL.TRUNK.NAME
      task.loss.momentum_teacher = build_model(
          config["MODEL"], config["OPTIMIZER"]
      )
      task.loss.momentum_teacher.to(task.device)
  
      if not is_fsdp_model:
          # Restore an hypothetical checkpoint
          if task.loss.checkpoint is not None:
              task.loss.load_state_dict(task.loss.checkpoint)
          # Else initialize from the student model
          else:
              task_model = get_no_ddp_model(task.model)
              teacher_model = get_no_ddp_model(
                  task.loss.momentum_teacher
              )
              teacher_model.load_state_dict(task_model.state_dict())
  
      # Setup SyncBN (useful for the XCiT)
      task.loss.momentum_teacher = (
          nn.SyncBatchNorm.convert_sync_batchnorm(
              task.loss.momentum_teacher
          )
      )
  
      # Wrap with DDP (needed for SyncBN) or FSDP
      if is_fsdp_model:
          fsdp_config = task.config["MODEL"]["FSDP_CONFIG"]
          task.loss.momentum_teacher = fsdp_wrapper(
              task.loss.momentum_teacher, **fsdp_config
          )
      else:
          task.loss.momentum_teacher = DistributedDataParallel(
              task.loss.momentum_teacher, device_ids=[task.device]
          )
  
      if is_fsdp_model:
          # Restore an hypothetical checkpoint
          if task.loss.checkpoint is not None:
              task.loss.load_state_dict(task.loss.checkpoint)
          else:
              # Else initialize from the student model
              task_model = task.base_model
              teacher_model = task.loss.momentum_teacher
              teacher_model.load_local_state_dict(
                  task_model.local_state_dict()
              )
          fsdp_recursive_reset_lazy_init(task.loss.momentum_teacher)
  
      # no gradients for teacher model
      for p in task.loss.momentum_teacher.parameters():
          p.requires_grad = False
  
file_path: vissl/vissl/hooks/hook_utils.py
imports: import logging
  import math
  import torch
  import torch.nn
  from classy_vision import tasks
  from torch.nn.parallel import DistributedDataParallel
  from vissl.config import AttrDict
  from vissl.models import build_model
  from vissl.models.model_helpers import get_no_ddp_model
  from vissl.utils.fsdp_utils import fsdp_recursive_reset_lazy_init, fsdp_wrapper
label_desc: Use image node_id 1 for calling the MomentumTeacherInLossHook obj's underlying member method code with example usage: obj.build_momentum_network(task, config) without return types