save grad_scaler state dict if use amp
Summary:
saving grad_scaler is better when resume training. Although the AMPTrainer implemented state_dict method,  in fact, the grad_scaler state_dict is not saved in checkpoints.
Pull Request resolved:
@@ -387,7 +387,6 @@ class DefaultTrainer(TrainerBase):
# Assume you want to save checkpoints together with logs/statistics
model,
cfg.OUTPUT_DIR,
- optimizer=optimizer,
trainer=weakref.proxy(self),
)
self.start_iter = 0
@@ -494,6 +493,15 @@ class DefaultTrainer(TrainerBase):
self._trainer.iter = self.iter
self._trainer.run_step()
+ def state_dict(self):
+ ret = super().state_dict()
+ ret["_trainer"] = self._trainer.state_dict()
+ return ret
+
+ def load_state_dict(self, state_dict):
+ super().load_state_dict(state_dict)
+ self._trainer.load_state_dict(state_dict["_trainer"])
+
@classmethod
def build_model(cls, cfg):
"""
