MLP_DIM in ViT is not adjustable through the config.

[Jeffkang-94]

Instructions To Reproduce the Issue:

Dino configuration contains the parameter named MLP_DIM which looks like being adjustable by the user. But actually, it is hard coded. See the line here

vissl/vissl/models/trunks/vision_transformer.py

Lines 149 to 155 in 4ca0ae3

	mlp_hidden_dim = int(dim * mlp_ratio)
	self.mlp = Mlp(
	in_features=dim,
	hidden_features=mlp_hidden_dim,
	act_layer=act_layer,
	drop=drop,
	)

I wanted to adjust the mlp_ratio or size of mlp_dim through the configuration as I've met several OOM issues because of the huge dataset,

  MODEL:
    TRUNK:
      NAME: vision_transformer
      VISION_TRANSFORMERS:
        IMAGE_SIZE: 512
        PATCH_SIZE: 16
        NUM_LAYERS: 3
        NUM_HEADS: 12
        HIDDEN_DIM: 192 
        # MLP_DIM: 768 # Not Used!
        CLASSIFIER: token
        DROPOUT_RATE: 0
        ATTENTION_DROPOUT_RATE: 0
        QKV_BIAS: True
        DROP_PATH_RATE: 0.1

But it seems that the mlp_dim is not adjustable through the configuration. Apart from it, it also raises some confusion for the user who wants to change the size of mlp dimension.

Expected behavior:

To be specific, with the configuration

MODEL:
    TRUNK:
      NAME: vision_transformer
      VISION_TRANSFORMERS:
        HIDDEN_DIM: 384
        MLP_DIM: 768 # Not Used!

mlp is supposed to have

mlp dim : Mlp(
  (fc1): Linear(in_features=384, out_features=768, bias=True)
  (act): GELU()
  (fc2): Linear(in_features=768, out_features=384, bias=True)
  (drop): Dropout(p=0, inplace=False)
)

but actually it has

mlp dim : Mlp(
  (fc1): Linear(in_features=384, out_features=1536, bias=True)
  (act): GELU()
  (fc2): Linear(in_features=1536, out_features=384, bias=True)
  (drop): Dropout(p=0, inplace=False)
)

because of the pre-defined value mlp_ratio=4.0 here

vissl/vissl/models/trunks/vision_transformer.py

Line 262 in 4ca0ae3

self.mlp_ratio = 4.0

[QuentinDuval]

Hello @Jeffkang-94,

Sorry for my delay in the response, and thanks a lot for raising the issue !

It should be solved at : 04788de

I also fixed the configurations that had the wrong MLP_DIM!

Thank you,
Quentin

[Jeffkang-94]

Thank you for your efforts! I think we can wrap-up this issue then.