class_docstr: Implementation of an AnyNet.
  
  See https://arxiv.org/abs/2003.13678 for details.
class_docstr_tok: ["Implementation","of","an","AnyNet",".","See","https",":","\/\/arxiv.org\/abs\/2003.13678","for","details","."]
code: def __init__(self, params: AnyNetParams):
      super().__init__()
  
      silu = None if get_torch_version() < [1, 7] else nn.SiLU()
      activation = {
          ActivationType.RELU: nn.ReLU(params.relu_in_place),
          ActivationType.SILU: silu,
      }[params.activation]
  
      if activation is None:
          raise RuntimeError(
              "SiLU activation is only supported since PyTorch 1.7"
          )
  
      # Ad hoc stem
      self.stem = {
          StemType.RES_STEM_CIFAR: ResStemCifar,
          StemType.RES_STEM_IN: ResStemIN,
          StemType.SIMPLE_STEM_IN: SimpleStemIN,
      }[params.stem_type](
          3,
          params.stem_width,
          params.bn_epsilon,
          params.bn_momentum,
          activation,
      )
  
      # Instantiate all the AnyNet blocks in the trunk
      block_fun = {
          BlockType.VANILLA_BLOCK: VanillaBlock,
          BlockType.RES_BASIC_BLOCK: ResBasicBlock,
          BlockType.RES_BOTTLENECK_BLOCK: ResBottleneckBlock,
          BlockType.RES_BOTTLENECK_LINEAR_BLOCK: ResBottleneckLinearBlock,
      }[params.block_type]
  
      current_width = params.stem_width
  
      self.trunk_depth = 0
  
      blocks = []
      for i, (
          width_out,
          stride,
          depth,
          group_width,
          bottleneck_multiplier,
      ) in enumerate(params.get_expanded_params()):
          blocks.append(
              (
                  f"block{i+1}",
                  AnyStage(
                      current_width,
                      width_out,
                      stride,
                      depth,
                      block_fun,
                      activation,
                      group_width,
                      bottleneck_multiplier,
                      params,
                      stage_index=i + 1,
                  ),
              )
          )
  
          self.trunk_depth += blocks[-1][1].stage_depth
  
          current_width = width_out
  
      self.trunk_output = nn.Sequential(OrderedDict(blocks))
  
      # Init weights and good to go
      self.init_weights()
  
imports: import math
  from collections import OrderedDict
  from enum import auto, Enum
  from typing import Any, Dict, Optional, Sequence
  import torch.nn
  from classy_vision.generic.util import get_torch_version
  from classy_vision.models import ClassyModel, register_model
  from classy_vision.models.squeeze_and_excitation_layer import SqueezeAndExcitationLayer
label_desc: Use image node_id 1 to create a new AnyNet object from inherited base classes: ClassyModel with example: obj = AnyNet(params)