|
- # 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.
- # ============================================================================
- """cifar_resnet50
- This sample code is applicable to Ascend.
- """
- import os
- import random
- import argparse
- from mindspore import dtype as mstype
- import mindspore.dataset as ds
- import mindspore.dataset.vision.c_transforms as C
- import mindspore.dataset.transforms.c_transforms as C2
- from mindspore.nn import SoftmaxCrossEntropyWithLogits
- from mindspore.communication import init
- from mindspore.nn import Momentum
- from mindspore import Model, context
- from mindspore.context import ParallelMode
- from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor
- from mindspore import load_checkpoint, load_param_into_net
- from resnet import resnet50
-
- def run(args):
- random.seed(1)
- parser = argparse.ArgumentParser(description='Image classification')
-
-
- context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target)
-
- if args.device_target == "Ascend":
- device_id = int(os.getenv('DEVICE_ID', '0'))
- context.set_context(device_id=device_id)
-
- if not args.do_eval and args.run_distribute:
- context.set_auto_parallel_context(device_num=args.device_num, parallel_mode=ParallelMode.DATA_PARALLEL,
- all_reduce_fusion_config=[140])
- init()
-
- epoch_size = args.epoch_size
- net = resnet50(args.batch_size, args.num_classes)
- ls = SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
- opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), 0.01, 0.9)
-
- model = Model(net, loss_fn=ls, optimizer=opt, metrics={'acc'})
-
- # as for train, users could use model.train
- if args.do_train:
- dataset = create_dataset(args)
- batch_num = dataset.get_dataset_size()
- config_ck = CheckpointConfig(save_checkpoint_steps=batch_num, keep_checkpoint_max=35)
- ckpoint_cb = ModelCheckpoint(prefix="train_resnet_cifar10", directory=args.train_url, config=config_ck)
- loss_cb = LossMonitor()
- model.train(epoch_size, dataset, callbacks=[ckpoint_cb, loss_cb])
-
- # as for evaluation, users could use model.eval
- if args.do_eval:
- if args.checkpoint_path:
- param_dict = load_checkpoint(args.checkpoint_path)
- load_param_into_net(net, param_dict)
- eval_dataset = create_dataset(args,training=False)
- res = model.eval(eval_dataset)
- print("result: ", res)
-
-
- def create_dataset(args,repeat_num=1, training=True):
- """
- create data for next use such as training or inferring
- """
- data_home = os.path.join(args.data_url,"cifar-10-batches-bin")
- assert os.path.exists(data_home), "the dataset path is invalid!"
- cifar_ds = ds.Cifar10Dataset(data_home)
-
- if args.run_distribute:
- rank_id = int(os.getenv('RANK_ID'))
- rank_size = int(os.getenv('RANK_SIZE'))
- cifar_ds = ds.Cifar10Dataset(data_home, num_shards=rank_size, shard_id=rank_id)
-
- resize_height = 224
- resize_width = 224
- rescale = 1.0 / 255.0
- shift = 0.0
-
- # define map operations
- random_crop_op = C.RandomCrop((32, 32), (4, 4, 4, 4)) # padding_mode default CONSTANT
- random_horizontal_op = C.RandomHorizontalFlip()
- resize_op = C.Resize((resize_height, resize_width)) # interpolation default BILINEAR
- rescale_op = C.Rescale(rescale, shift)
- normalize_op = C.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
- changeswap_op = C.HWC2CHW()
- type_cast_op = C2.TypeCast(mstype.int32)
-
- c_trans = []
- if training:
- c_trans = [random_crop_op, random_horizontal_op]
- c_trans += [resize_op, rescale_op, normalize_op,
- changeswap_op]
-
- # apply map operations on images
- cifar_ds = cifar_ds.map(operations=type_cast_op, input_columns="label")
- cifar_ds = cifar_ds.map(operations=c_trans, input_columns="image")
-
- # apply shuffle operations
- cifar_ds = cifar_ds.shuffle(buffer_size=10)
-
- # apply batch operations
- cifar_ds = cifar_ds.batch(batch_size=args.batch_size, drop_remainder=True)
-
- # apply repeat operations
- cifar_ds = cifar_ds.repeat(repeat_num)
-
- return cifar_ds
|