text: # Copyright 2019 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.
  # ============================================================================
  
  import numpy as np
  import os
  import pytest
  
  import mindspore.context as context
  import mindspore.dataset as ds
  import mindspore.dataset.transforms.c_transforms as C
  import mindspore.dataset.transforms.vision.c_transforms as CV
  import mindspore.nn as nn
  from mindspore import Tensor
  from mindspore.common import dtype as mstype
  from mindspore.common.initializer import initializer
  from mindspore.dataset.transforms.vision import Inter
  from mindspore.model_zoo.lenet import LeNet5
  from mindspore.nn import Dense, TrainOneStepCell, WithLossCell
  from mindspore.nn.metrics import Accuracy
  from mindspore.nn.optim import Momentum
  from mindspore.ops import operations as P
  from mindspore.train import Model
  from mindspore.train.callback import LossMonitor
  
  context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
  
  
  class LeNet(nn.Cell):
      def __init__(self):
          super(LeNet, self).__init__()
          self.relu = P.ReLU()
          self.batch_size = 1
          weight1 = Tensor(np.ones([6, 3, 5, 5]).astype(np.float32) * 0.01)
          weight2 = Tensor(np.ones([16, 6, 5, 5]).astype(np.float32) * 0.01)
          self.conv1 = nn.Conv2d(3, 6, (5, 5), weight_init=weight1, stride=1, padding=0, pad_mode='valid')
          self.conv2 = nn.Conv2d(6, 16, (5, 5), weight_init=weight2, pad_mode='valid', stride=1, padding=0)
          self.pool = nn.MaxPool2d(kernel_size=2, stride=2, pad_mode="valid")
  
          self.reshape = P.Reshape()
          self.reshape1 = P.Reshape()
  
          self.fc1 = Dense(400, 120)
          self.fc2 = Dense(120, 84)
          self.fc3 = Dense(84, 10)
  
      def construct(self, input_x):
          output = self.conv1(input_x)
          output = self.relu(output)
          output = self.pool(output)
          output = self.conv2(output)
          output = self.relu(output)
          output = self.pool(output)
          output = self.reshape(output, (self.batch_size, -1))
          output = self.fc1(output)
          output = self.fc2(output)
          output = self.fc3(output)
          return output
  
  
  def multisteplr(total_steps, gap, base_lr=0.9, gamma=0.1, dtype=mstype.float32):
      lr = []
      for step in range(total_steps):
          lr_ = base_lr * gamma ** (step // gap)
          lr.append(lr_)
      return Tensor(np.array(lr), dtype)
  
  
  @pytest.mark.level0
  @pytest.mark.platform_x86_gpu_training
  @pytest.mark.env_onecard
  def test_train_lenet():
      epoch = 100
      net = LeNet()
      momentum = initializer(Tensor(np.array([0.9]).astype(np.float32)), [1])
      learning_rate = multisteplr(epoch, 30)
  
      optimizer = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), learning_rate, momentum)
      criterion = nn.SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=True)
      net_with_criterion = WithLossCell(net, criterion)
      train_network = TrainOneStepCell(net_with_criterion, optimizer)  # optimizer
      train_network.set_train()
      losses = []
      for i in range(epoch):
          data = Tensor(np.ones([net.batch_size, 3, 32, 32]).astype(np.float32) * 0.01)
          label = Tensor(np.ones([net.batch_size]).astype(np.int32))
          loss = train_network(data, label)
          losses.append(loss)
      print(losses)
  
  
  def create_dataset(data_path, batch_size=32, repeat_size=1,
                     num_parallel_workers=1):
      """
      create dataset for train or test
      """
      # define dataset
      mnist_ds = ds.MnistDataset(data_path)
  
      resize_height, resize_width = 32, 32
      rescale = 1.0 / 255.0
      shift = 0.0
      rescale_nml = 1 / 0.3081
      shift_nml = -1 * 0.1307 / 0.3081
  
      # define map operations
      resize_op = CV.Resize((resize_height, resize_width), interpolation=Inter.LINEAR)  # Bilinear mode
      rescale_nml_op = CV.Rescale(rescale_nml, shift_nml)
      rescale_op = CV.Rescale(rescale, shift)
      hwc2chw_op = CV.HWC2CHW()
      type_cast_op = C.TypeCast(mstype.int32)
  
      # apply map operations on images
      mnist_ds = mnist_ds.map(input_columns="label", operations=type_cast_op, num_parallel_workers=num_parallel_workers)
      mnist_ds = mnist_ds.map(input_columns="image", operations=resize_op, num_parallel_workers=num_parallel_workers)
      mnist_ds = mnist_ds.map(input_columns="image", operations=rescale_op, num_parallel_workers=num_parallel_workers)
      mnist_ds = mnist_ds.map(input_columns="image", operations=rescale_nml_op, num_parallel_workers=num_parallel_workers)
      mnist_ds = mnist_ds.map(input_columns="image", operations=hwc2chw_op, num_parallel_workers=num_parallel_workers)
  
      # apply DatasetOps
      buffer_size = 10000
      mnist_ds = mnist_ds.shuffle(buffer_size=buffer_size)  # 10000 as in LeNet train script
      mnist_ds = mnist_ds.batch(batch_size, drop_remainder=True)
      mnist_ds = mnist_ds.repeat(repeat_size)
  
      return mnist_ds
  
  
  @pytest.mark.level0
  @pytest.mark.platform_x86_gpu_training
  @pytest.mark.env_onecard
  def test_train_and_eval_lenet():
      context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
      network = LeNet5(10)
      net_loss = nn.SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=True, reduction="mean")
      net_opt = nn.Momentum(network.trainable_params(), 0.01, 0.9)
      model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()})
  
      print("============== Starting Training ==============")
      ds_train = create_dataset(os.path.join('/home/workspace/mindspore_dataset/mnist', "train"), 32, 1)
      model.train(1, ds_train, callbacks=[LossMonitor()], dataset_sink_mode=True)
  
      print("============== Starting Testing ==============")
      ds_eval = create_dataset(os.path.join('/home/workspace/mindspore_dataset/mnist', "test"), 32, 1)
      acc = model.eval(ds_eval, dataset_sink_mode=True)
      print("============== Accuracy:{} ==============".format(acc))
  