class_docstr: timm model adapter
  # FIXME this adapter is a work in progress, may change in ways that break weight compat
docstr: lock modules
  Args:
      unlocked_groups (int): leave last n layer groups unlocked (default: 0)
docstr_tok: ["lock","modules","Args",":","unlocked_groups","(","int",")",":","leave","last","n","layer","groups","unlocked","(","default",":","0",")"]
code: def lock(self, unlocked_groups=0, freeze_bn_stats=False):
      """lock modules
      Args:
          unlocked_groups (int): leave last n layer groups unlocked (default: 0)
      """
      if not unlocked_groups:
          # lock full model
          for param in self.trunk.parameters():
              param.requires_grad = False
          if freeze_bn_stats:
              freeze_batch_norm_2d(self.trunk)
      else:
          # NOTE: partial freeze requires latest timm (master) branch and is subject to change
          try:
              # FIXME import here until API stable and in an official release
              from timm.models.helpers import (
                  group_modules,
                  group_parameters,
              )
          except ImportError:
              raise RuntimeError(
                  "Please install latest timm `pip install git+https://github.com/rwightman/pytorch-image-models`"
              )
          matcher = self.trunk.group_matcher()
          gparams = group_parameters(self.trunk, matcher)
          max_layer_id = max(gparams.keys())
          max_layer_id = max_layer_id - unlocked_groups
          for group_idx in range(max_layer_id + 1):
              group = gparams[group_idx]
              for param in group:
                  self.trunk.get_parameter(param).requires_grad = False
          if freeze_bn_stats:
              gmodules = group_modules(
                  self.trunk, matcher, reverse=True
              )
              gmodules = {
                  k for k, v in gmodules.items() if v <= max_layer_id
              }
              freeze_batch_norm_2d(self.trunk, gmodules)
  
imports: import math
  import warnings
  from collections import OrderedDict
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn
  from torch import nn
  from lavis.models.clip_models.utils import freeze_batch_norm_2d
label_desc: Use image node_id 2 for calling the TimmModel obj's underlying member method code with example usage: obj.lock(unlocked_groups, freeze_bn_stats) without return types