diff --git a/ppcls/engine/engine.py b/ppcls/engine/engine.py index e468b29a3..741dd36d1 100755 --- a/ppcls/engine/engine.py +++ b/ppcls/engine/engine.py @@ -242,61 +242,8 @@ class Engine(object): self.config["Optimizer"], self.config["Global"]["epochs"], self.iter_per_epoch // self.update_freq, [self.model, self.train_loss_func]) - - # AMP training and evaluating - self.amp = "AMP" in self.config and self.config["AMP"] is not None - self.amp_eval = False - # for amp - if self.amp: - AMP_RELATED_FLAGS_SETTING = {'FLAGS_max_inplace_grad_add': 8, } - if paddle.is_compiled_with_cuda(): - AMP_RELATED_FLAGS_SETTING.update({ - 'FLAGS_cudnn_batchnorm_spatial_persistent': 1 - }) - paddle.set_flags(AMP_RELATED_FLAGS_SETTING) - - self.scale_loss = self.config["AMP"].get("scale_loss", 1.0) - self.use_dynamic_loss_scaling = self.config["AMP"].get( - "use_dynamic_loss_scaling", False) - self.scaler = paddle.amp.GradScaler( - init_loss_scaling=self.scale_loss, - use_dynamic_loss_scaling=self.use_dynamic_loss_scaling) - - self.amp_level = self.config['AMP'].get("level", "O1") - if self.amp_level not in ["O1", "O2"]: - msg = "[Parameter Error]: The optimize level of AMP only support 'O1' and 'O2'. The level has been set 'O1'." - logger.warning(msg) - self.config['AMP']["level"] = "O1" - self.amp_level = "O1" - - self.amp_eval = self.config["AMP"].get("use_fp16_test", False) - # TODO(gaotingquan): Paddle not yet support FP32 evaluation when training with AMPO2 - if self.mode == "train" and self.config["Global"].get( - "eval_during_train", - True) and self.amp_level == "O2" and self.amp_eval == False: - msg = "PaddlePaddle only support FP16 evaluation when training with AMP O2 now. " - logger.warning(msg) - self.config["AMP"]["use_fp16_test"] = True - self.amp_eval = True - - if self.mode == "train": - self.model, self.optimizer = paddle.amp.decorate( - models=self.model, - optimizers=self.optimizer, - level=self.amp_level, - save_dtype='float32') - elif self.amp_eval: - self.model = paddle.amp.decorate( - models=self.model, - level=self.amp_level, - save_dtype='float32') - - if self.mode == "train" and len(self.train_loss_func.parameters( - )) > 0: - self.train_loss_func = paddle.amp.decorate( - models=self.train_loss_func, - level=self.amp_level, - save_dtype='float32') + # amp + self._init_amp() # build EMA model self.ema = "EMA" in self.config and self.mode == "train" @@ -513,7 +460,9 @@ class Engine(object): batch_tensor = paddle.to_tensor(batch_data) if self.amp and self.amp_eval: - with paddle.amp.auto_cast(level=self.amp_level): + with paddle.amp.auto_cast( + level=self.amp_level, + use_promote=self.use_promote): out = self.model(batch_tensor) else: out = self.model(batch_tensor) @@ -578,6 +527,62 @@ class Engine(object): f"Export succeeded! The inference model exported has been saved in \"{self.config['Global']['save_inference_dir']}\"." ) + def _init_amp(self): + self.amp = "AMP" in self.config and self.config["AMP"] is not None + self.amp_eval = False + + if self.amp: + AMP_RELATED_FLAGS_SETTING = {'FLAGS_max_inplace_grad_add': 8, } + if paddle.is_compiled_with_cuda(): + AMP_RELATED_FLAGS_SETTING.update({ + 'FLAGS_cudnn_batchnorm_spatial_persistent': 1 + }) + paddle.set_flags(AMP_RELATED_FLAGS_SETTING) + + self.scale_loss = self.config["AMP"].get("scale_loss", 1.0) + self.use_dynamic_loss_scaling = self.config["AMP"].get( + "use_dynamic_loss_scaling", False) + self.scaler = paddle.amp.GradScaler( + init_loss_scaling=self.scale_loss, + use_dynamic_loss_scaling=self.use_dynamic_loss_scaling) + + self.use_promote = self.config['AMP'].get("use_promote", False) + self.amp_level = self.config['AMP'].get("level", "O1") + if self.amp_level not in ["O1", "O2"]: + msg = "[Parameter Error]: The optimize level of AMP only support 'O1' and 'O2'. The level has been set 'O1'." + logger.warning(msg) + self.config['AMP']["level"] = "O1" + self.amp_level = "O1" + + self.amp_eval = self.config["AMP"].get("use_fp16_test", False) + # TODO(gaotingquan): Paddle not yet support FP32 evaluation when training with AMPO2 + if self.mode == "train" and self.config["Global"].get( + "eval_during_train", + True) and self.amp_level == "O2" and self.amp_eval == False: + msg = "PaddlePaddle only support FP16 evaluation when training with AMP O2 now. " + logger.warning(msg) + self.config["AMP"]["use_fp16_test"] = True + self.amp_eval = True + + if self.mode == "train": + self.model, self.optimizer = paddle.amp.decorate( + models=self.model, + optimizers=self.optimizer, + level=self.amp_level, + save_dtype='float32') + elif self.amp_eval: + self.model = paddle.amp.decorate( + models=self.model, + level=self.amp_level, + save_dtype='float32') + + if self.mode == "train" and len(self.train_loss_func.parameters( + )) > 0: + self.train_loss_func = paddle.amp.decorate( + models=self.train_loss_func, + level=self.amp_level, + save_dtype='float32') + class ExportModel(TheseusLayer): """