class_docstr: Support TIMM image backbones.
  Refer to https://github.com/rwightman/pytorch-image-models
docstr: Load a pretrained image backbone from TIMM.
  
  Parameters
  ----------
  prefix
      The prefix of the TimmAutoModelForImagePrediction model.
  checkpoint_name
      Name of the timm checkpoint, or local parent directory of the saved finetuned timm weights and config.
  num_classes
      The number of classes. 1 for a regression task.
  mix_choice
      Choice used for mixing multiple images. We now support.
      - all_images
          The images are directly averaged and passed to the model.
      - all_logits
          The logits output from individual images are averaged to generate the final output.
  pretrained
      Whether using the pretrained timm models. If pretrained=True, download the pretrained model.
docstr_tok: ["Load","a","pretrained","image","backbone","from","TIMM",".","Parameters","--","--","--","--","--","prefix","The","prefix","of","the","TimmAutoModelForImagePrediction","model",".","checkpoint_name","Name","of","the","timm","checkpoint",",","or","local","parent","directory","of","the","saved","finetuned","timm","weights","and","config",".","num_classes","The","number","of","classes",".","1","for","a","regression","task",".","mix_choice","Choice","used","for","mixing","multiple","images",".","We","now","support",".","-","all_images","The","images","are","directly","averaged","and","passed","to","the","model",".","-","all_logits","The","logits","output","from","individual","images","are","averaged","to","generate","the","final","output",".","pretrained","Whether","using","the","pretrained","timm","models",".","If","pretrained=True",",","download","the","pretrained","model","."]
code: def __init__(
      self,
      prefix: str,
      checkpoint_name: str,
      num_classes: Optional[int] = 0,
      mix_choice: Optional[str] = "all_logits",
      pretrained: Optional[bool] = True,
  ):
      """
      Load a pretrained image backbone from TIMM.
  
      Parameters
      ----------
      prefix
          The prefix of the TimmAutoModelForImagePrediction model.
      checkpoint_name
          Name of the timm checkpoint, or local parent directory of the saved finetuned timm weights and config.
      num_classes
          The number of classes. 1 for a regression task.
      mix_choice
          Choice used for mixing multiple images. We now support.
          - all_images
              The images are directly averaged and passed to the model.
          - all_logits
              The logits output from individual images are averaged to generate the final output.
      pretrained
          Whether using the pretrained timm models. If pretrained=True, download the pretrained model.
      """
      super().__init__()
      # In TIMM, if num_classes==0, then create_model would automatically set self.model.head = nn.Identity()
      logger.debug(f"initializing {checkpoint_name}")
      if os.path.exists(checkpoint_name):
          checkpoint_path = f"{checkpoint_name}/pytorch_model.bin"
          try:
              with open(f"{checkpoint_name}/config.json") as f:
                  self.config = json.load(f)
                  pretrained_cfg = self.config.get("pretrained_cfg", {})
                  for k, v in pretrained_cfg.items():
                      if k not in self.config:
                          self.config[k] = v
                  self.checkpoint_name = self.config.get(
                      "architecture", None
                  )
                  self.model = create_model(
                      self.checkpoint_name,
                      checkpoint_path=checkpoint_path,
                      num_classes=0,
                  )
                  # create a head with new num_classes
                  self.head = (
                      Linear(
                          in_features=self.config["num_features"],
                          out_features=num_classes,
                      )
                      if num_classes > 0
                      else nn.Identity()
                  )
                  self.num_classes = (
                      num_classes if num_classes is not None else 0
                  )
          except:
              raise ValueError(
                  f"Timm model path {checkpoint_name} does not exist or model is invalid."
              )
      else:
          self.checkpoint_name = checkpoint_name
          self.model = create_model(
              checkpoint_name,
              pretrained=pretrained,
              num_classes=num_classes,
          )
          self.head = get_model_head(model=self.model)
          self.config = self.model.default_cfg
          self.num_classes = self.model.num_classes
  
      self.pretrained = pretrained
      self.out_features = self.model.num_features
      self.global_pool = (
          self.model.global_pool
          if hasattr(self.model, "global_pool")
          else None
      )
      self.model.reset_classifier(0)  # remove the internal head
  
      self.mix_choice = mix_choice
      logger.debug(f"mix_choice: {mix_choice}")
  
      self.prefix = prefix
  
      self.name_to_id = self.get_layer_ids()
      self.head_layer_names = [
          n for n, layer_id in self.name_to_id.items() if layer_id == 0
      ]
  
imports: import json
  import logging
  import os
  from typing import Dict, List, Optional
  import torch
  from timm import create_model
  from timm.layers.linear import Linear
  from torch import nn
  from ..constants import AUTOMM, COLUMN, COLUMN_FEATURES, FEATURES, IMAGE, IMAGE_VALID_NUM, LABEL, LOGITS, MASKS
  from .utils import assign_layer_ids, get_column_features, get_model_head
label_desc: Use image node_id 1 to create a new TimmAutoModelForImagePrediction object from inherited base classes: nn with example: obj = TimmAutoModelForImagePrediction(prefix, checkpoint_name, num_classes, mix_choice, pretrained)