• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2020 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15import os
16import random
17import argparse
18import numpy as np
19from resnet import resnet50
20
21import mindspore.common.dtype as mstype
22import mindspore.ops.functional as F
23from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor, TimeMonitor
24from mindspore.train.serialization import load_checkpoint, load_param_into_net
25import mindspore.dataset as ds
26import mindspore.dataset.transforms.c_transforms as C
27import mindspore.dataset.vision.c_transforms as vision
28import mindspore.nn as nn
29from mindspore import Tensor
30from mindspore import context
31from mindspore.communication.management import init
32from mindspore.nn.optim.momentum import Momentum
33from mindspore.ops import operations as P
34from mindspore.train.model import Model
35from mindspore.context import ParallelMode
36
37random.seed(1)
38np.random.seed(1)
39ds.config.set_seed(1)
40
41parser = argparse.ArgumentParser(description='Image classification')
42parser.add_argument('--run_distribute', type=bool, default=False, help='Run distribute')
43parser.add_argument('--device_num', type=int, default=1, help='Device num.')
44parser.add_argument('--do_train', type=bool, default=True, help='Do train or not.')
45parser.add_argument('--do_eval', type=bool, default=False, help='Do eval or not.')
46parser.add_argument('--epoch_size', type=int, default=1, help='Epoch size.')
47parser.add_argument('--batch_size', type=int, default=32, help='Batch size.')
48parser.add_argument('--num_classes', type=int, default=10, help='Num classes.')
49parser.add_argument('--checkpoint_path', type=str, default=None, help='Checkpoint file path')
50parser.add_argument('--dataset_path', type=str, default="/var/log/npu/datasets/cifar", help='Dataset path')
51args_opt = parser.parse_args()
52
53device_id = int(os.getenv('DEVICE_ID'))
54
55data_home = args_opt.dataset_path
56
57context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
58context.set_context(device_id=device_id)
59
60
61def create_dataset(repeat_num=1, training=True):
62    data_dir = data_home + "/cifar-10-batches-bin"
63    if not training:
64        data_dir = data_home + "/cifar-10-verify-bin"
65    data_set = ds.Cifar10Dataset(data_dir, num_samples=32)
66
67    if args_opt.run_distribute:
68        rank_id = int(os.getenv('RANK_ID'))
69        rank_size = int(os.getenv('RANK_SIZE'))
70        data_set = ds.Cifar10Dataset(data_dir, num_shards=rank_size, shard_id=rank_id, num_samples=32)
71
72    resize_height = 224
73    resize_width = 224
74    rescale = 1.0 / 255.0
75    shift = 0.0
76
77    # define map operations
78    random_crop_op = vision.RandomCrop((32, 32), (4, 4, 4, 4))  # padding_mode default CONSTANT
79    random_horizontal_op = vision.RandomHorizontalFlip()
80    resize_op = vision.Resize((resize_height, resize_width))  # interpolation default BILINEAR
81    rescale_op = vision.Rescale(rescale, shift)
82    normalize_op = vision.Normalize((0.4465, 0.4822, 0.4914), (0.2010, 0.1994, 0.2023))
83    changeswap_op = vision.HWC2CHW()
84    type_cast_op = C.TypeCast(mstype.int32)
85
86    c_trans = []
87    if training:
88        c_trans = [random_crop_op, random_horizontal_op]
89    c_trans += [resize_op, rescale_op, normalize_op,
90                changeswap_op]
91
92    # apply map operations on images
93    data_set = data_set.map(operations=type_cast_op, input_columns="label")
94    data_set = data_set.map(operations=c_trans, input_columns="image")
95
96    # apply repeat operations
97    data_set = data_set.repeat(repeat_num)
98
99    # apply shuffle operations
100    data_set = data_set.shuffle(buffer_size=10)
101
102    # apply batch operations
103    data_set = data_set.batch(batch_size=args_opt.batch_size, drop_remainder=True)
104
105    return data_set
106
107
108class CrossEntropyLoss(nn.Cell):
109    def __init__(self):
110        super(CrossEntropyLoss, self).__init__()
111        self.cross_entropy = P.SoftmaxCrossEntropyWithLogits()
112        self.mean = P.ReduceMean()
113        self.one_hot = P.OneHot()
114        self.one = Tensor(1.0, mstype.float32)
115        self.zero = Tensor(0.0, mstype.float32)
116
117    def construct(self, logits, label):
118        label = self.one_hot(label, F.shape(logits)[1], self.one, self.zero)
119        loss_func = self.cross_entropy(logits, label)[0]
120        loss_func = self.mean(loss_func, (-1,))
121        return loss_func
122
123
124if __name__ == '__main__':
125    if not args_opt.do_eval and args_opt.run_distribute:
126        context.set_auto_parallel_context(device_num=args_opt.device_num, parallel_mode=ParallelMode.DATA_PARALLEL,
127                                          all_reduce_fusion_config=[140])
128        init()
129
130    context.set_context(mode=context.GRAPH_MODE)
131    epoch_size = args_opt.epoch_size
132    net = resnet50(args_opt.batch_size, args_opt.num_classes)
133    loss = CrossEntropyLoss()
134    opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), 0.01, 0.9)
135
136    model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc'})
137
138    if args_opt.do_train:
139        dataset = create_dataset(1)
140        batch_num = dataset.get_dataset_size()
141        config_ck = CheckpointConfig(save_checkpoint_steps=batch_num * 5, keep_checkpoint_max=10)
142        ckpoint_cb = ModelCheckpoint(prefix="train_resnet_cifar10", directory="./", config=config_ck)
143        time_cb = TimeMonitor(data_size=batch_num)
144        loss_cb = LossMonitor()
145        model.train(epoch_size, dataset, callbacks=[ckpoint_cb, loss_cb, time_cb])
146
147    if args_opt.do_eval:
148        if args_opt.checkpoint_path:
149            param_dict = load_checkpoint(args_opt.checkpoint_path)
150            load_param_into_net(net, param_dict)
151        net.set_train(False)
152        eval_dataset = create_dataset(1, training=False)
153        res = model.eval(eval_dataset)
154        print("result: ", res)
155