Adding multi-layer perceptron in ops

[datumbox]

We should avoid using ViT's MLP block from Swin:

vision/torchvision/models/swin_transformer.py

Line 15 in 9d9cfab

from .vision_transformer import MLPBlock

The specific layer is very common and has been previously requested at #4333

This PR:

• Adds a generic MLP block to TorchVision.
• Handles BC on ViT.
• Replaces MLPBlock with MLP on Swin and patches the unreleased weights (weights uploaded on S3 and manifold).

References:

• TorchVision
• Multimodal
• DETR
• Timm
• TF

Proof that the new API doesn't break Swin. The minor differences are expected due to the known non-deterministic behaviour of some kernels:

srun -p dev --cpus-per-task=96 -t 24:00:00 --gpus-per-node=1 torchrun --nproc_per_node=1 train.py --model swin_t --test-only -b 1 --weights="Swin_T_Weights.IMAGENET1K_V1"
Test:  Acc@1 81.476 Acc@5 95.780

srun -p dev --cpus-per-task=96 -t 24:00:00 --gpus-per-node=1 torchrun --nproc_per_node=1 train.py --model swin_s --test-only -b 1 --weights="Swin_S_Weights.IMAGENET1K_V1"
Test:  Acc@1 83.182 Acc@5 96.366

srun -p dev --cpus-per-task=96 -t 24:00:00 --gpus-per-node=1 torchrun --nproc_per_node=1 train.py --model swin_b --test-only -b 1 --weights="Swin_B_Weights.IMAGENET1K_V1"
Test:  Acc@1 83.584 Acc@5 96.636

cc @ankitade

[NicolasHug]

Annotations in docstring make me sad :'(

[datumbox]

I know how you feel above this. :( I think all the callables all over TorchVision are added like that to provide info on what they are supposed to return.

[NicolasHug]

API LGTM, thanks @datumbox

[jdsgomes]

LGTM, feel free to merge after the same changes are done in Swin. If you prefer I can have a second look once thats done

[datumbox]

This is a straight move from convnext to ops (no weight patching needed) but to avoid doing everything on a single PR I plan to do it on a follow up.

[thomasbbrunner]

It is not very clear for me why there's a Dropout layer after the last layer. I saw that it was present in the previous MLPBlock class, but no other implementation of MLP with dropout (that I could find) has a dropout layer on the output. Including the one in the multimodal package.

Maybe this was something specific for the usecase of MLPBlock? If so, this should not be in this class.

[datumbox]

You are right there are various implementations of MLP, some of which don't have at all dropout, some have at the middle but not at the end or some have everywhere. If you check the references, you will see that all patterns exist. Our implementation is like that because it replaces MLP layers used in existing models like ViT and Swin. We also try to support more complex variations with more than 2 linear layers. Your observation is correct though that if one wanted to avoid having dropout at the end, the current implementation wouldn't let them. Since that variant is also valid, perhaps it's worth making this update in a non-BC way with a new boolean that controls the appearance of Dropout at the end or not. WDYT?

[thomasbbrunner]

I think your suggestion sounds very nice. What would be the default value be for the boolean? I guess that setting it to True (with dropout) would cause no breaking changes. At the same time, I would say that not having a dropout in the last layer is more common (default) configuration? Also, I'd be intested in working on this, whichever option is chosen.

[datumbox]

I guess that setting it to True (with dropout) would cause no breaking changes.

Yes, you are right we will need to maintain BC. Note that using True is the "default" setup on TorchVision at the moment as literally all existing models require dropout everywhere.

I'd be intested in working on this

Sounds great, let me recommend the following. Could you start an issue, summarizing what you said here and providing a few references of the usage of MLP with a middle dropout but without the final one? Providing a few examples from real-world vision architectures will help build a stronger case. Once we clarify the details on the issue, we can discuss a potential PR. 😃

[thomasbbrunner]

Ok! I am a bit short on time at the moment, but will have more time in the upcoming weeks. Nevertheless, I'm interested in this and will be working on it!