|
-
-
- import os
- import argparse
-
- import moxing as mox
- from config import mnist_cfg as cfg
- from dataset import create_dataset
- from lenet import LeNet5
- import json
- import mindspore.nn as nn
- from mindspore import context
- from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor, TimeMonitor
- from mindspore.train import Model
- from mindspore.nn.metrics import Accuracy
- from mindspore import load_checkpoint, load_param_into_net
- from mindspore.context import ParallelMode
- from mindspore.communication.management import init, get_rank
- import time
-
-
- parser = argparse.ArgumentParser(description='MindSpore Lenet Example')
- ### --multi_data_url,--train_url,--device_target,These 3 parameters must be defined first in a multi-dataset,
- ### otherwise an error will be reported.
- ### There is no need to add these parameters to the running parameters of the Qizhi platform,
- ### because they are predefined in the background, you only need to define them in your code.
-
- parser.add_argument('--multi_data_url',
- help='path to multi dataset',
- default= '/cache/data/')
-
- parser.add_argument('--ckpt_url',
- help='pre_train_model path in obs')
-
- # parser.add_argument('--model_url',
- # help='pre_train_model path in obs')
-
- parser.add_argument('--train_url',
- help='model folder to save/load',
- default= '/cache/output/')
-
- parser.add_argument(
- '--device_target',
- type=str,
- default="Ascend",
- choices=['Ascend', 'CPU'],
- help='device where the code will be implemented (default: Ascend),if to use the CPU on the Qizhi platform:device_target=CPU')
-
- parser.add_argument('--epoch_size',
- type=int,
- default=5,
- help='Training epochs.')
-
- if __name__ == "__main__":
- args, unknown = parser.parse_known_args()
- data_dir = '../'
- train_dir = '../output'
- model_dir = '../'
- if not os.path.exists(data_dir):
- os.makedirs(data_dir)
- if not os.path.exists(train_dir):
- os.makedirs(train_dir)
- if not os.path.exists(model_dir):
- os.makedirs(model_dir)
- ds_train = create_dataset(os.path.join(data_dir, "train"), cfg.batch_size)
- network = LeNet5(cfg.num_classes)
- net_loss = nn.SoftmaxCrossEntropyWithLogits(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())
- load_param_into_net(network, load_checkpoint(os.path.join(model_dir, "checkpoint.ckpt")))
- if args.device_target != "Ascend":
- model = Model(network,net_loss,net_opt,metrics={"accuracy": Accuracy()})
- else:
- model = Model(network, net_loss,net_opt,metrics={"accuracy": Accuracy()},amp_level="O2")
- config_ck = CheckpointConfig(save_checkpoint_steps=cfg.save_checkpoint_steps,
- keep_checkpoint_max=cfg.keep_checkpoint_max)
- #Note that this method saves the model file on each card. You need to specify the save path on each card.
- # In this example, get_rank() is added to distinguish different paths.
-
- outputDirectory = train_dir
- ckpoint_cb = ModelCheckpoint(prefix="checkpoint_lenet",
- directory=outputDirectory,
- config=config_ck)
- print("============== Starting Training ==============")
- epoch_size = cfg['epoch_size']
- if (args.epoch_size):
- epoch_size = args.epoch_size
- print('epoch_size is: ', epoch_size)
- model.train(epoch_size,
- ds_train,
- callbacks=[time_cb, ckpoint_cb,
- LossMonitor()])
|