Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

MLP as used in Vision Transformer, MLP-Mixer and related networks

Classifier head w/ configurable global pooling and dropout.

r"""Applies a linear transformation to the incoming data: :math:`y = xA^T + b` Wraps torch.nn.Linear to support AMP + torchscript usage by manual

Image to Patch Embedding. Different with ViT use 1 conv layer, we use 4 conv layers to do patch embedding

Selectable global pooling layer with dynamic input kernel size

Downscaling with maxpool then double conv

Downscaling with maxpool then double conv

Upscaling then double conv

Upscaling then double conv

A downsample pooling module supporting several maxpool and avgpool modes * 'max' - MaxPool2d w/ kernel_size 3, stride 2, padding 1 * 'max2' -

A downsample pooling module supporting several maxpool and avgpool modes * 'max' - MaxPool2d w/ kernel_size 3, stride 2, padding 1 * 'max2' -

LayerNorm w/ fast norm option

Pre-activation (v2) bottleneck block.

Pre-activation (v2) bottleneck block.

MLP using 1x1 convs that keeps spatial dims

(convolution => [BN] => ReLU) * 2

(convolution => [BN] => ReLU) * 2

Convolutional Position Encoding. Note: This module is similar to the conditional position encoding in CPVT.

Convolutional relative position encoding.

Build Trunk Blocks This ended up being somewhat of a cross between https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mna

Factorized attention with convolutional relative position encoding class.

Image to Patch Embedding

Serial block class. Note: In this implementation, each serial block only contains a conv-attention and a FFN (MLP) module.

Multi-Head Attention

Channels-last multi-head attention (B, ..., C)

Channels-last multi-head attention (B, ..., C)

TNT Block