• 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.context as context
23import mindspore.dataset as de
24import mindspore.dataset.transforms.c_transforms as C
25import mindspore.dataset.vision.c_transforms as vision
26import mindspore.nn as nn
27from mindspore import Tensor
28from mindspore.communication.management import init
29from mindspore.nn.optim.momentum import Momentum
30from mindspore.ops import functional as F
31from mindspore.ops import operations as P
32from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor
33from mindspore.train.model import Model
34from mindspore.context import ParallelMode
35
36random.seed(1)
37np.random.seed(1)
38de.config.set_seed(1)
39
40parser = argparse.ArgumentParser(description='Image classification')
41parser.add_argument('--run_distribute', type=bool, default=False, help='Run distribute')
42parser.add_argument('--device_num', type=int, default=1, help='Device num.')
43parser.add_argument('--do_train', type=bool, default=True, help='Do train or not.')
44parser.add_argument('--do_eval', type=bool, default=False, help='Do eval or not.')
45parser.add_argument('--epoch_size', type=int, default=1, help='Epoch size.')
46parser.add_argument('--batch_size', type=int, default=4, help='Batch size.')
47parser.add_argument('--num_classes', type=int, default=10, help='Num classes.')
48parser.add_argument('--checkpoint_path', type=str, default=None, help='Checkpoint file path')
49parser.add_argument('--dataset_path', type=str, default="/var/log/npu/datasets/cifar", help='Dataset path')
50args_opt = parser.parse_args()
51
52device_id = int(os.getenv('DEVICE_ID'))
53
54data_home = args_opt.dataset_path
55
56context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
57context.set_context(device_id=device_id)
58
59
60def create_dataset(repeat_num=1, training=True):
61    data_dir = data_home + "/cifar-10-batches-bin"
62    if not training:
63        data_dir = data_home + "/cifar-10-verify-bin"
64    ds = de.Cifar10Dataset(data_dir)
65
66    if args_opt.run_distribute:
67        rank_id = int(os.getenv('RANK_ID'))
68        rank_size = int(os.getenv('RANK_SIZE'))
69        ds = de.Cifar10Dataset(data_dir, num_shards=rank_size, shard_id=rank_id)
70
71    resize_height = 224
72    resize_width = 224
73    rescale = 1.0 / 255.0
74    shift = 0.0
75
76    # define map operations
77    random_crop_op = vision.RandomCrop((32, 32), (4, 4, 4, 4))  # padding_mode default CONSTANT
78    random_horizontal_op = vision.RandomHorizontalFlip()
79    resize_op = vision.Resize((resize_height, resize_width))  # interpolation default BILINEAR
80    rescale_op = vision.Rescale(rescale, shift)
81    normalize_op = vision.Normalize((0.4465, 0.4822, 0.4914), (0.2010, 0.1994, 0.2023))
82    changeswap_op = vision.HWC2CHW()
83    type_cast_op = C.TypeCast(mstype.int32)
84
85    c_trans = []
86    if training:
87        c_trans = [random_crop_op, random_horizontal_op]
88    c_trans += [resize_op, rescale_op, normalize_op,
89                changeswap_op]
90
91    # apply map operations on images
92    ds = ds.map(operations=type_cast_op, input_columns="label")
93    ds = ds.map(operations=c_trans, input_columns="image")
94
95    # apply repeat operations
96    ds = ds.repeat(repeat_num)
97
98    # apply shuffle operations
99    ds = ds.shuffle(buffer_size=10)
100
101    # apply batch operations
102    ds = ds.batch(batch_size=args_opt.batch_size, drop_remainder=True)
103
104    return ds
105
106
107class CrossEntropyLoss(nn.Cell):
108    def __init__(self):
109        super(CrossEntropyLoss, self).__init__()
110        self.cross_entropy = P.SoftmaxCrossEntropyWithLogits()
111        self.mean = P.ReduceMean()
112        self.one_hot = P.OneHot()
113        self.one = Tensor(1.0, mstype.float32)
114        self.zero = Tensor(0.0, mstype.float32)
115
116    def construct(self, logits, label):
117        label = self.one_hot(label, F.shape(logits)[1], self.one, self.zero)
118        loss_func = self.cross_entropy(logits, label)[0]
119        loss_func = self.mean(loss_func, (-1,))
120        return loss_func
121
122
123if __name__ == '__main__':
124    if not args_opt.do_eval and args_opt.run_distribute:
125        context.set_auto_parallel_context(device_num=args_opt.device_num, parallel_mode=ParallelMode.DATA_PARALLEL)
126        context.set_auto_parallel_context(all_reduce_fusion_split_indices=[140])
127        init()
128
129    context.set_context(mode=context.GRAPH_MODE)
130    epoch_size = args_opt.epoch_size
131    net = resnet50(args_opt.batch_size, args_opt.num_classes)
132    loss = CrossEntropyLoss()
133    opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), 0.01, 0.9)
134
135    model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc'})
136
137    if args_opt.do_train:
138        dataset = create_dataset(epoch_size)
139        batch_num = dataset.get_dataset_size()
140        config_ck = CheckpointConfig(save_checkpoint_steps=batch_num * 5, keep_checkpoint_max=10)
141        ckpoint_cb = ModelCheckpoint(prefix="train_resnet_cifar10", directory="./", config=config_ck)
142        loss_cb = LossMonitor()
143        model.train(epoch_size, dataset, callbacks=[ckpoint_cb, loss_cb])
144
145    if args_opt.do_eval:
146        eval_dataset = create_dataset(1, training=False)
147        res = model.eval(eval_dataset)
148        print("result: ", res)
149    checker = os.path.exists("./normal_memreuse.ir")
150    assert checker, True
151