Torch AMP + ShardedDDP (#667)

Summary: Pull Request resolved: #667 Adding mixed precision support in conjunction with Classy Reviewed By: mannatsingh Differential Revision: D25400577 fbshipit-source-id: 687eefac5cb4ea24d89ac22a92492389d136f015
facebookresearch · Dec 11, 2020 · a17e50a · a17e50a
1 parent ff37fea
commit a17e50a
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/classy_vision/tasks/classification_task.py b/classy_vision/tasks/classification_task.py
@@ -455,6 +455,10 @@ def set_amp_args(self, amp_args: Optional[Dict[str, Any]]):
                     "Apex AMP is required but Apex is not installed, cannot enable AMP"
                 )
 
+            # Set Torch AMP grad scaler, used to prevent gradient underflow
+            elif self.amp_type == AmpType.PYTORCH:
+                self.amp_grad_scaler = TorchGradScaler()
+
             logging.info(f"AMP enabled with args {amp_args}")
         return self
 
@@ -744,8 +748,6 @@ def prepare(self):
                     self.base_model, self.optimizer.optimizer = apex.amp.initialize(
                         self.base_model, self.optimizer.optimizer, **self.amp_args
                     )
-            elif self.amp_type == AmpType.PYTORCH:
-                self.amp_grad_scaler = TorchGradScaler()
 
         if self.simulated_global_batchsize is not None:
             if self.simulated_global_batchsize % self.get_global_batchsize() != 0: