!9 modify alexnet

Merge pull request !9 from wukesong/modify-alexnet

chapter04/alexnet/generator_lr.py

+# 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

chapter04/alexnet/main.py

@@ -17,14 +17,16 @@ AlexNet example tutorial
 Usage:
      python alexnet.py
 with --device_target=GPU: After 20 epoch training, the accuracy is up to 80%
-with --device_target=Ascend: After 10 epoch training, the accuracy is up to 81%
+with --device_target=Ascend: After 30 epoch training, the accuracy is up to 88%
 """
 
 import argparse
 from config import alexnet_cfg as cfg
 from alexnet import AlexNet
+from generator_lr import get_lr
 import mindspore.dataset as ds
 import mindspore.nn as nn
+from mindspore import Tensor
 from mindspore import context
 from mindspore.train.serialization import load_checkpoint, load_param_into_net
 from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor
@@ -75,7 +77,7 @@ if __name__ == "__main__":
     parser.add_argument('--data_path', type=str, default="./", help='path where the dataset is saved')
     parser.add_argument('--ckpt_path', type=str, default="./ckpt", help='if mode is test, must provide\
                         path where the trained ckpt file')
-    parser.add_argument('--dataset_sink_mode', type=bool, default=False, help='dataset_sink_mode is False or True')
+    parser.add_argument('--dataset_sink_mode', type=bool, default=True, help='dataset_sink_mode is False or True')
     args = parser.parse_args()
 
     context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target)
@@ -83,7 +85,9 @@ if __name__ == "__main__":
     network = AlexNet(cfg.num_classes)
     loss = nn.SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=True, reduction="mean")
     repeat_size = cfg.epoch_size
-    opt = nn.Momentum(network.trainable_params(), cfg.learning_rate, cfg.momentum)
+    # when batch_size=32, steps is 1562
+    lr = Tensor(get_lr(0, cfg.learning_rate, cfg.epoch_size, 1562))
+    opt = nn.Momentum(network.trainable_params(), lr, cfg.momentum)
     model = Model(network, loss, opt, metrics={"Accuracy": Accuracy()})  # test
 
     if args.mode == 'train':