Skip to content

ConvNeXt#

eqxvision.models.ConvNeXt #

A simple port of torchvision.models.convnext.

__init__(self, block_setting: Sequence[_CNBlockConfig], stochastic_depth_prob: float = 0.0, layer_scale: float = 1e-06, num_classes: int = 1000, block: Optional[eqx.Module] = None, norm_layer: Optional[eqx.Module] = None, *, key: jax.random.PRNGKey = None) #
  • block_setting: Configuration of the computational blocks
  • stochastic_depth_prob: Probability of dropping a sample along channels
  • layer_scale: Scale applied to the output of computational stem
  • num_classes: Number of classes in the classification task. Also controls the final output shape (num_classes,)
  • block: The block type used within the network
  • norm_layer: Normalisation applied to the intermediate outputs
  • key: A jax.random.PRNGKey used to provide randomness for parameter initialisation. (Keyword only argument.)
__call__(self, x: Array, *, 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.convnext_tiny(*, torch_weights: str = None, **kwargs: Any) -> ConvNeXt #

ConvNeXt Tiny model architecture from the "A ConvNet for the 2020s" <https://arxiv.org/abs/2201.03545>_ paper.

Arguments:

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

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

ConvNeXt Small model architecture from the A ConvNet for the 2020s paper.

Arguments:

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

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

ConvNeXt Base model architecture from the A ConvNet for the 2020s paper.

Arguments:

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

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

ConvNeXt Large model architecture from the A ConvNet for the 2020s paper.

Arguments:

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