Squeeze-Excitation

eqxvision.layers.SqueezeExcitation

A simple port of torchvision.ops.misc.SqueezeExcitation

__init__(self, input_channels: int, squeeze_channels: int, activation: Callable = None, scale_activation: Callable = None, *, key: jax.random.PRNGKey = None)

Arguments:

• input_channels: Number of channels in the input image
• squeeze_channels: Number of squeeze channels
• activation: delta activation. Defaults to jax.nn.relu
• scale_activation: sigma`` activation. Defaults tojax.nn.sigmoid`
• 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: Forwarded to individual eqx.Module attributes