Skip to content

MobileNet-V2#

eqxvision.models.MobileNetV2 #

A simple port of torchvision.models.mobilenetv2

__init__(self, num_classes: int = 1000, width_mult: float = 1.0, inverted_residual_setting: Optional[List[List[int]]] = None, round_nearest: int = 8, block: Optional[eqx.Module] = None, norm_layer: Optional[eqx.Module] = None, dropout: float = 0.2, *, key: Optional[jax.random.PRNGKey] = None) #

Arguments:

  • num_classes: Number of classes in the classification task. Also controls the final output shape (num_classes,). Defaults to 1000
  • width_mult: Adjusts number of channels in each layer by this amount inverted_residual_setting: Network structure
  • inverted_residual_setting: Network structure
  • round_nearest: Round the number of channels in each layer to be a multiple of this number Set to 1 to turn off rounding
  • block: Module specifying inverted residual building block for mobilenet
  • norm_layer: Module specifying the normalization layer to use
  • dropout: The dropout probability
  • key: A jax.random.PRNGKey used to provide randomness for parameter initialisation. (Keyword only argument.)
__call__(self, x, *, key: jax.random.PRNGKey) -> Array #

Arguments:

  • x: The input JAX array
  • key: Required parameter. Utilised by few layers such as Dropout or DropPath

eqxvision.models.mobilenet_v2(torch_weights: str = None, **kwargs: Any) -> MobileNetV2 #

Constructs a MobileNetV2 architecture from MobileNetV2: Inverted Residuals and Linear Bottlenecks.

Arguments:

  • torch_weights: A Path or URL for the PyTorch weights. Defaults to None