Use detectron2 implementation for Attention layer in ViTDet.

Author: hgaiser

references/detection/train.py

@@ -256,6 +256,25 @@ def main(args):
         )
     elif opt_name == "adamw":
         optimizer = torch.optim.AdamW(parameters, lr=args.lr, weight_decay=args.weight_decay)
+    elif opt_name == "vitdet":
+        from torchvision.models.vision_transformer import get_default_optimizer_params, get_vit_lr_decay_rate
+        from functools import partial
+
+        optimizer = torch.optim.AdamW(
+            params=get_default_optimizer_params(
+                model,
+                # params.model is meant to be set to the model object, before instantiating
+                # the optimizer.
+                base_lr=args.lr,
+                weight_decay_norm=0.0,
+                # TODO: Adjust num_layers for specific model. Currently this assumes ViT-B.
+                lr_factor_func=partial(get_vit_lr_decay_rate, num_layers=12, lr_decay_rate=0.7),
+                overrides={"pos_embed": {"weight_decay": 0.0}},
+            ),
+            lr=args.lr,
+            betas=(0.9, 0.999),
+            weight_decay=0.1,
+        )
     else:
         raise RuntimeError(f"Invalid optimizer {args.opt}. Only SGD and AdamW are supported.")
 

torchvision/models/detection/mask_rcnn.py

@@ -672,7 +672,7 @@ def maskrcnn_vit_b_16_sfpn(
     elif num_classes is None:
         num_classes = 91
 
-    backbone = vit_b_16(weights=weights_backbone, progress=progress, include_head=False)
+    backbone = vit_b_16(weights=weights_backbone, progress=progress, include_head=False, image_size=1024)
     backbone = _vit_sfpn_extractor(backbone)
     model = MaskRCNN(backbone, num_classes=num_classes, **kwargs)