Skip to content

ShuffleNet-V2#

eqxvision.models.ShuffleNetV2 #

A simple port of torchvision.models.shufflenetv2

__init__(self, stages_repeats: List[int], stages_out_channels: List[int], num_classes: int = 1000, inverted_residual: eqx.Module = <class 'eqxvision.models.classification.shufflenetv2._InvertedResidual'>, *, key: Optional[jax.random.PRNGKey] = None) #

Arguments:

  • stages_repeats: Number of times a block is repeated for each stage
  • stages_out_channels: Output at each stage
  • num_classes: Number of classes in the classification task. Also controls the final output shape (num_classes,). Defaults to 1000
  • inverted_residual: Network structure
  • key: A jax.random.PRNGKey used to provide randomness for parameter initialisation. (Keyword only argument.)
__call__(self, x, *, key: Optional[jax.random.PRNGKey] = None) -> Array #

Arguments:

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

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

Constructs a ShuffleNetV2 with 0.5x output channels, as described in ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design.

Arguments:

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

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

Constructs a ShuffleNetV2 with 1.0x output channels, as described in ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design.

Arguments:

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

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

Constructs a ShuffleNetV2 with 1.5x output channels, as described in ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design.

Arguments:

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

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

Constructs a ShuffleNetV2 with 2.0x output channels, as described in ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design.

Arguments:

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