support amp training for detection models

[xiaohu2015]

The pr is about #4509.
Since amp is supported on classification training, I also modify some files to support amp training on detetction models.

cc @datumbox

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit bfc9225 (more details on the Dr. CI page):

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

[datumbox]

Thanks for the PR @xiaohu2015!

There are some related linter issues, but overall it looks good.

@prabhat00155 could you have also a look since you did the classification one?

[xiaohu2015]

@datumbox Thanks. I have test the amp code, it works well.

[datumbox]

@xiaohu2015 There are a couple of more linter issues (spaces). Have a look a the CI job, at the end it will show you the errors. For your convenience here are the things you need to change to keep it happy:

diff --git a/references/detection/train.py b/references/detection/train.py
index 5c50dcfa..ae13a32b 100644
--- a/references/detection/train.py
+++ b/references/detection/train.py
@@ -143,7 +143,7 @@ def get_args_parser(add_help=True):
 
     # Prototype models only
     parser.add_argument("--weights", default=None, type=str, help="the weights enum name to load")
-    
+
     # Mixed precision training parameters
     parser.add_argument("--amp", action="store_true", help="Use torch.cuda.amp for mixed precision training")
 
@@ -211,9 +211,9 @@ def main(args):
 
     params = [p for p in model.parameters() if p.requires_grad]
     optimizer = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
-    
+
     scaler = torch.cuda.amp.GradScaler() if args.amp else None
-    
+
     args.lr_scheduler = args.lr_scheduler.lower()
     if args.lr_scheduler == "multisteplr":
         lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.lr_steps, gamma=args.lr_gamma)

Other than that, the changes look good to me. :)

I'll leave it to @prabhat00155 to do the final checks on our side and merge when ready.

[prabhat00155]

Thanks @xiaohu2015!

[github-actions]

Hey @prabhat00155!

You merged this PR, but no labels were added. The list of valid labels is available at https://github.com/pytorch/vision/blob/main/.github/process_commit.py