From f8fa51002f330fb62d9cd32072ad469367cb5e2c Mon Sep 17 00:00:00 2001 From: VectorSL Date: Thu, 15 Oct 2020 16:53:05 +0800 Subject: [PATCH] gpu modelzoo lenet is not suitable for amp --- model_zoo/official/cv/lenet/train.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/model_zoo/official/cv/lenet/train.py b/model_zoo/official/cv/lenet/train.py index 05a5e72155e..fa7c6534d80 100644 --- a/model_zoo/official/cv/lenet/train.py +++ b/model_zoo/official/cv/lenet/train.py @@ -58,7 +58,10 @@ if __name__ == "__main__": config_ck = CheckpointConfig(save_checkpoint_steps=cfg.save_checkpoint_steps, keep_checkpoint_max=cfg.keep_checkpoint_max) ckpoint_cb = ModelCheckpoint(prefix="checkpoint_lenet", directory=args.ckpt_path, config=config_ck) - model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()}, amp_level="O2") + if args.device_target != "Ascend": + model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()}) + else: + model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()}, amp_level="O2") print("============== Starting Training ==============") model.train(cfg['epoch_size'], ds_train, callbacks=[time_cb, ckpoint_cb, LossMonitor()],