MODULE_08_PYTORCH_PAPER_REPLICATING_SECTION_ 6.2_Replicating_Equation_3_with_PyTorch_layers

[Krishanu2206]

SECTION DOUBT LINK : https://www.learnpytorch.io/08_pytorch_paper_replicating/#62-replicating-equation-3-with-pytorch-layers

Q - Can I use the MLP module given in the torchvision.ops documentation directly here? MLP :
https://pytorch.org/vision/main/generated/torchvision.ops.MLP.html#torchvision.ops.MLP

I am attaching my code here :

CODE :

from torchvision import ops
from torch import nn

MLP_block = torchvision.ops.MLP(in_channels=768, 
                                    hidden_channels=[3072, 768], 
                                    norm_layer=nn.LayerNorm,
                                    activation_layer=nn.GELU, 
                                    dropout=0.1 ) ## dropout as specified in table-3

MLP_block

OUTPUT :

MLP(
  (0): Linear(in_features=768, out_features=3072, bias=True)
  (1): LayerNorm((3072,), eps=1e-05, elementwise_affine=True)
  (2): GELU(approximate='none')
  (3): Dropout(p=0.1, inplace=False)
  (4): Linear(in_features=3072, out_features=768, bias=True)
  (5): Dropout(p=0.1, inplace=False)
)

CODE :

# Pass output of MSABlock through DirectMLPLayer
patched_image_through_mlp_block = MLP_block(patched_image_through_msa_block)
print(f"Input shape of MLP block: {patched_image_through_msa_block.shape}")
print(f"Output shape MLP block: {patched_image_through_mlp_block.shape}")

OUTPUT :

Input shape of MLP block: torch.Size([1, 197, 768])
Output shape MLP block: torch.Size([1, 197, 768])

The only difference in the MLP module and the MLP module in the video is the placement of the torch.nn.LayerNorm() .