OpenI_Cloudbrain_Example/npu_mnist_example/train_npu.py

"""
示例选用的数据集是MNISTData.zip
数据集结构是：
 MNISTData.zip
  ├── test
  │   ├── t10k-images-idx3-ubyte
  │   └── t10k-labels-idx1-ubyte
  └── train
      ├── train-images-idx3-ubyte
      └── train-labels-idx1-ubyte 
      
使用注意事项：
1、在代码中加入args, unknown = parser.parse_known_args()，可忽略掉--ckpt_url参数报错等参数问题
2、用户需要调用openi的python sdk包
"""

import os
import argparse
from config import mnist_cfg as cfg
from dataset import create_dataset
from dataset_distributed import create_dataset_parallel
from lenet import LeNet5
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.context import ParallelMode
from mindspore.communication.management import init, get_rank
import time
#导入openi包
from openi.context import prepare, upload_openi


parser = argparse.ArgumentParser(description='MindSpore Lenet Example')
                   
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()，可忽略掉--ckpt_url参数报错等参数问题
    args, unknown = parser.parse_known_args()
    data_dir = ''  
    pretrain_dir = ''
    train_dir = ''

    #回传结果到openi
    upload_openi()        
  
    device_num = int(os.getenv('RANK_SIZE'))
    #使用单卡时
    if device_num == 1:
        context.set_context(mode=context.GRAPH_MODE,device_target=args.device_target)
        #初始化导入数据集和预训练模型到容器内
        openi_context = prepare()
        data_dir = openi_context.dataset_path
        pretrain_dir = openi_context.pretrain_model_path
        train_dir = openi_context.output_path
        #使用数据集的方式  
        ds_train = create_dataset(os.path.join(data_dir + "/MNISTData", "train"),  cfg.batch_size)
    #使用多卡时        
    if device_num > 1:
        # set device_id and init for multi-card training
        context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=int(os.getenv('ASCEND_DEVICE_ID')))
        context.reset_auto_parallel_context()
        context.set_auto_parallel_context(device_num = device_num, parallel_mode=ParallelMode.DATA_PARALLEL, gradients_mean=True, parameter_broadcast=True)
        init()
        #Copying obs data does not need to be executed multiple times, just let the 0th card copy the data
        local_rank=int(os.getenv('RANK_ID'))
        if local_rank%8==0:
            ###初始化导入数据集和预训练模型到容器内
            openi_context = prepare()
            #初始化导入数据集和预训练模型到容器内
            openi_context = prepare()
            data_dir = openi_context.dataset_path
            pretrain_dir = openi_context.pretrain_model_path
            train_dir = openi_context.output_path
            #Set a cache file to determine whether the data has been copied to obs. 
            #If this file exists during multi-card training, there is no need to copy the dataset multiple times.
            f = open("/cache/download_input.txt", 'w')    
            f.close()
            try:
                if os.path.exists("/cache/download_input.txt"):
                    print("download_input succeed")
            except Exception as e:
                print("download_input failed")
        while not os.path.exists("/cache/download_input.txt"):
            time.sleep(1)               
        ds_train = create_dataset_parallel(os.path.join(data_dir + "/MNISTData", "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())

    if args.device_target != "Ascend":
        model = Model(network,
                      net_loss,
                      net_opt,
                      metrics={"accuracy"})
    else:
        model = Model(network,
                      net_loss,
                      net_opt,
                      metrics={"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.
    if device_num == 1:
        outputDirectory = train_dir + "/"
    if device_num > 1:
        outputDirectory = train_dir + "/" + str(get_rank()) + "/"
    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()])

    ###上传训练结果到启智平台,注意必须将要输出的模型存储在openi_context.output_path   
    upload_openi()