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# ============================================================================ 15 16import os 17 18import numpy as np 19 20import mindspore.communication.management as distributedTool 21import mindspore.nn as nn 22from mindspore import context 23from mindspore.nn.metrics import Accuracy 24from mindspore.train import Model 25from mindspore.train.callback import LossMonitor, TimeMonitor 26from tests.models.official.cv.lenet.src.dataset import create_dataset 27from tests.models.official.cv.lenet.src.lenet import LeNet5 28 29np.set_printoptions(threshold=np.inf) 30device_num = 2 31device_id = int(os.getenv('DEVICE_ID')) 32rank_id = 0 33 34 35def setup_module(): 36 global device_num 37 global rank_id 38 context.set_context(mode=context.GRAPH_MODE, device_target="Ascend") 39 context.set_context(device_id=device_id) 40 distributedTool.init() 41 rank_id = distributedTool.get_rank() 42 device_num = distributedTool.get_group_size() 43 context.set_auto_parallel_context(device_num=device_num, global_rank=device_id, parameter_broadcast=True) 44 45 46def teardown_module(): 47 distributedTool.release() 48 49 50def test_all_trains(): 51 ds_train = create_dataset(os.path.join('/home/workspace/mindspore_dataset/mnist', "train"), 32, 1) 52 53 network = LeNet5(10) 54 net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean") 55 net_opt = nn.Momentum(network.trainable_params(), 0.01, 0.9) 56 time_cb = TimeMonitor(data_size=ds_train.get_dataset_size()) 57 58 model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()}) 59 60 print("============== Starting Training ==============") 61 model.train(1, ds_train, callbacks=[time_cb, LossMonitor()]) 62