Add Video SwinTransformer

[oke-aditya]

Closes #6499

I think ready for a initial round of review. Will love to understand a bit more and then continue.

I will mostly to wrap up this before the end of month. Little bit more work is pending 😄

[oke-aditya]

I got the hang of SwinTransformer2d. Soon I will take this up.

[oke-aditya]

This was mega learning and refresher for me as I haven't at all worked on models / attention.
(My ML skills probably are waning as I don't much work on it these days)

OK so after a long wait, I have completely read both the papers swin image and swin video. I'm currently working to align the image and video codebases and would then write down the model configurations. I have a few comments though

[oke-aditya]

I think couple of things should be added.

1. Logging.
2. torch.fx.wrap()

I would love to understand what torch.fx.wrap() does. 😄

[oke-aditya]

I think there is no such function in 2d case. This is just in 3d

[YosuaMichael]

Just FYI, there is an equivalent in these lines: https://github.com/pytorch/vision/blob/main/torchvision/models/swin_transformer.py#L193-L207

But yeah, we dont wrap them as separate function in 2d implementation

[oke-aditya]

I think we can refactor the 2d implementation a bit, the 3d looks much better here 😄

[oke-aditya]

And this too can be rewritten

[YosuaMichael]

Yeah, for V2 we should separate this into separate function

[oke-aditya]

Somehow I am not convinced with this style of writing both in 2d and 3d, why can't we give defaults in the function definition? Is there any reason?

[YosuaMichael]

Actually I follow the 2d version style for this, I guess for the previous one it kinda make sense if we use partial(...), but I agree that we might as well put this as default directly in the function definition for patch_embed. Can't really think the advantage using this style now.

[oke-aditya]

I hope we can refactor the implementation, but it might be BC breaking. (minor bc break though)

[oke-aditya]

Just curious what was equivalent to this in 2d case and where is the codeblock for same?

[YosuaMichael]

This actually implemented directly in the main class in the 2d version:
https://github.com/pytorch/vision/blob/main/torchvision/models/swin_transformer.py#L548-L557

In the 3d version we need to separate this into separate class because when Omnivore use this swin3d as encoder, it need to use different way for creating patch embedding: https://github.com/facebookresearch/multimodal/blob/main/torchmultimodal/models/omnivore.py#L144

[oke-aditya]

I wish we can separate in the 2d version too. It will make the paper implementation bit more clear as the block diagram illustrated.

[YosuaMichael]

Hi @oke-aditya , nice work! For this first review, I address some of your question and concern. I will do further review tomorrow :)

[YosuaMichael]

Just FYI, there is an equivalent in these lines: https://github.com/pytorch/vision/blob/main/torchvision/models/swin_transformer.py#L193-L207

But yeah, we dont wrap them as separate function in 2d implementation

[YosuaMichael]

I think the only way to remove the loop is to do comparison one by one (3 if statement).

Btw, take note that on this specific logic, it is actually difference between video swin transformer and the 2d.
On our 2d implementation, we dont change the window size but only set the shift to be 0.

On the original 3d implementation, we actually resize the window_size like how it is implemented here.

[YosuaMichael]

Yeah, for V2 we should separate this into separate function

[YosuaMichael]

This actually implemented directly in the main class in the 2d version:
https://github.com/pytorch/vision/blob/main/torchvision/models/swin_transformer.py#L548-L557

In the 3d version we need to separate this into separate class because when Omnivore use this swin3d as encoder, it need to use different way for creating patch embedding: https://github.com/facebookresearch/multimodal/blob/main/torchmultimodal/models/omnivore.py#L144

[YosuaMichael]

maybe we should use nn.Identity to make mypy happy?

[YosuaMichael]

Actually I follow the 2d version style for this, I guess for the previous one it kinda make sense if we use partial(...), but I agree that we might as well put this as default directly in the function definition for patch_embed. Can't really think the advantage using this style now.

[oke-aditya]

Ok so this seems almost done. I guess we are ready to port weights. Can you verify the configurations @YosuaMichael ?

[oke-aditya]

Just a quick update.
We plan to continue working on this after the release. This will give us sufficient time to port all the weights!

cc @datumbox

[oke-aditya]

@datumbox
Quick update, @YosuaMichael and I have resumed work on this. @YosuaMichael might kick off a run to check for initial set of weights. Hopefully we are in luck.

[YosuaMichael]

Thanks @oke-aditya for the ported weight. I modify the code from this branch a bit by adding the weight and transforms to mimic the original paper. Then I ran the weight to do testing on Kinetics-400 validation data using the following command:

python -u ~/script/run_with_submitit.py \
    --timeout 3000 --ngpus 8 --nodes 1 \
    --data-path="/datasets01/kinetics/070618/400/" \
    --batch-size=16 --test-only \
    --clip-len 32 --frame-rate 15 --clips-per-video 12 \
    --cache-dataset \
    --model {model_name} --weights="{weight_name}"

Note: I use clips-per-video==12 to simulate the paper that use 4 spatial x 3 temporal clips since as of now we dont have the spatial clips capability. Thanks @datumbox suggestion for this.

Where the pair of the model_name and weight_name with their corresponding accuracy are as follow:

1. model_name=swin3d_t, weight_name=Swin3D_T_Weights.KINETICS400_V1

 * Video Acc@1 77.715 Video Acc@5 93.519

2. model_name=swin3d_s, weight_name=Swin3D_S_Weights.KINETICS400_V1

 * Video Acc@1 79.521 Video Acc@5 94.158

3. model_name=swin3d_b, weight_name=Swin3D_B_Weights.KINETICS400_V1

 * Video Acc@1 79.427 Video Acc@5 94.386

4. model_name=swin3d_b, weight_name=Swin3D_B_Weights.IN22K_KINETICS400_V1

 * Video Acc@1 81.643 Video Acc@5 95.574

[datumbox]

@YosuaMichael Looks good. We should add the info about the input parameters passed to the ref script on the docs similar to what we do for MViT:

vision/torchvision/models/video/mvit.py

Lines 616 to 619 in cffb7f7

	"_docs": (
	"The weights were ported from the paper. The accuracies are estimated on video-level "
	"with parameters `frame_rate=7.5`, `clips_per_video=5`, and `clip_len=16`"
	),

[oke-aditya]

Nice, I will wrap this over the weekend!

[oke-aditya]

I kind of disagree to the name IN22K_KINETICS400_V1 Since for Swag models we put the names for other pretrained dataset as last.
E.g
Default one

 IMAGENET1K_V2 = Weights(

Swag one

IMAGENET1K_SWAG_E2E_V1 = Weights(

I would suggest for KINETICS we do

KINETICS400_V1 

and

KINETICS400_IMAGENET22K_V1 = Weights(

Edit:

Also I suggest for consistency the default weights for swin_b should be on IMAGENET_1K rather than IMAGENET_22K

[oke-aditya]

Btw also how do I get these 3 fields?

 "min_size": (224, 224),
 "min_temporal_size": 16,
"num_params": 36610672,

[datumbox]

@oke-aditya Thanks for the update.

KINETICS400_IMAGENET22K_V1

Sounds good to me.

For min_size its the minimum spatial size that can be sent to your architecture without throwing error. Sometimes maxpooling or conv layers impose a minimum size below which the network returns an error. I usually try a couple of values until I find the one that breaks the architecture and then set the min_size to the next available permitted value.

For min_temporal_size it's exactly the same but for the temporal T dimension.

Finally for num_params, just do sum(p.numel() for p in model.parameters()) to estimate them.

[YosuaMichael]

Looks good @oke-aditya , I agree with the naming KINETICS_IMAGENET22K_V1!
For min_size, temporal_size, and num_params it is as @datumbox explained.
I think swin transformer 3d dont really have min_size and min_temporal_size since we do padding in PatchEmbed3d (this means the min_size=(1,1) and min_temporal_size=1), however please give it a try first.

[oke-aditya]

So I checked the min_size and temporal size. It seems even (1,1) would work fine. And 1 for temporal size is fine too

>>> from torchvision.models.video import swin3d_b
>>> v_tensor = torch.randn(1, 3, 1, 1, 1) # b c t h w
>>> v_tensor = v_tensor.cuda()
>>> model = swin3d_b()
>>> out = model(v_tensor)

[oke-aditya]

@YosuaMichael @datumbox I have added the FLOPs and weights size by running the script in #6936

I think we should be good to merge.

[YosuaMichael]

@oke-aditya there seems lint error about the typing https://app.circleci.com/pipelines/github/pytorch/vision/21954/workflows/1179cdc7-f7db-4700-87c7-e45ef42af46c/jobs/1777451
could you resolve this?

[oke-aditya]

This should be good to go. CI failure is unrelated

[YosuaMichael]

Thanks for the hard work @oke-aditya, it looks good now!

[YosuaMichael]

The test failures are unrelated so we merged.