diff --git a/example/alexnet_cifar10/generator_lr.py b/example/alexnet_cifar10/generator_lr.py new file mode 100755 index 00000000000..1856124d3ef --- /dev/null +++ b/example/alexnet_cifar10/generator_lr.py @@ -0,0 +1,44 @@ +# Copyright 2020 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""learning rate generator""" +import numpy as np + + +def get_lr(current_step, lr_max, total_epochs, steps_per_epoch): + """ + generate learning rate array + + Args: + current_step(int): current steps of the training + lr_max(float): max learning rate + total_epochs(int): total epoch of training + steps_per_epoch(int): steps of one epoch + + Returns: + np.array, learning rate array + """ + lr_each_step = [] + total_steps = steps_per_epoch * total_epochs + decay_epoch_index = [0.8 * total_steps] + for i in range(total_steps): + if i < decay_epoch_index[0]: + lr = lr_max + else: + lr = lr_max * 0.1 + lr_each_step.append(lr) + lr_each_step = np.array(lr_each_step).astype(np.float32) + learning_rate = lr_each_step[current_step:] + + return learning_rate diff --git a/example/alexnet_cifar10/train.py b/example/alexnet_cifar10/train.py index 622df2d4042..0a288ea1db1 100644 --- a/example/alexnet_cifar10/train.py +++ b/example/alexnet_cifar10/train.py @@ -21,12 +21,14 @@ python train.py --data_path /YourDataPath import argparse from config import alexnet_cfg as cfg from dataset import create_dataset +from generator_lr import get_lr import mindspore.nn as nn from mindspore import context +from mindspore import Tensor from mindspore.train import Model from mindspore.nn.metrics import Accuracy from mindspore.model_zoo.alexnet import AlexNet -from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor +from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor, TimeMonitor if __name__ == "__main__": @@ -43,16 +45,17 @@ if __name__ == "__main__": network = AlexNet(cfg.num_classes) loss = nn.SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=True, reduction="mean") - opt = nn.Momentum(network.trainable_params(), cfg.learning_rate, cfg.momentum) + lr = Tensor(get_lr(0, cfg.learning_rate, cfg.epoch_size, cfg.save_checkpoint_steps)) + opt = nn.Momentum(network.trainable_params(), lr, cfg.momentum) model = Model(network, loss, opt, metrics={"Accuracy": Accuracy()}) # test print("============== Starting Training ==============") ds_train = create_dataset(args.data_path, cfg.batch_size, - cfg.epoch_size, - "train") + cfg.epoch_size) + time_cb = TimeMonitor(data_size=ds_train.get_dataset_size()) config_ck = CheckpointConfig(save_checkpoint_steps=cfg.save_checkpoint_steps, keep_checkpoint_max=cfg.keep_checkpoint_max) ckpoint_cb = ModelCheckpoint(prefix="checkpoint_alexnet", directory=args.ckpt_path, config=config_ck) - model.train(cfg.epoch_size, ds_train, callbacks=[ckpoint_cb, LossMonitor()], + model.train(cfg.epoch_size, ds_train, callbacks=[time_cb, ckpoint_cb, LossMonitor()], dataset_sink_mode=args.dataset_sink_mode) diff --git a/example/lenet_mnist/train.py b/example/lenet_mnist/train.py index d58d1a101b4..6c5cfb6fa2a 100644 --- a/example/lenet_mnist/train.py +++ b/example/lenet_mnist/train.py @@ -25,7 +25,7 @@ from dataset import create_dataset import mindspore.nn as nn from mindspore.model_zoo.lenet import LeNet5 from mindspore import context -from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor +from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor, TimeMonitor from mindspore.train import Model from mindspore.nn.metrics import Accuracy @@ -40,19 +40,20 @@ if __name__ == "__main__": args = parser.parse_args() - context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target) + context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, enable_mem_reuse=False) + ds_train = create_dataset(os.path.join(args.data_path, "train"), + cfg.batch_size, + cfg.epoch_size) network = LeNet5(cfg.num_classes) net_loss = nn.SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=True, reduction="mean") net_opt = nn.Momentum(network.trainable_params(), cfg.lr, cfg.momentum) + time_cb = TimeMonitor(data_size=ds_train.get_dataset_size()) config_ck = CheckpointConfig(save_checkpoint_steps=cfg.save_checkpoint_steps, keep_checkpoint_max=cfg.keep_checkpoint_max) ckpoint_cb = ModelCheckpoint(prefix="checkpoint_lenet", config=config_ck) model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()}) - ds_train = create_dataset(os.path.join(args.data_path, "train"), - cfg.batch_size, - cfg.epoch_size) print("============== Starting Training ==============") - model.train(cfg['epoch_size'], ds_train, callbacks=[ckpoint_cb, LossMonitor()], + model.train(cfg['epoch_size'], ds_train, callbacks=[time_cb, ckpoint_cb, LossMonitor()], dataset_sink_mode=args.dataset_sink_mode)