• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2019 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
19import pytest
20
21import mindspore.context as context
22import mindspore.dataset as ds
23import mindspore.dataset.transforms.c_transforms as C
24import mindspore.dataset.vision.c_transforms as CV
25import mindspore.nn as nn
26from mindspore import Tensor, ParameterTuple
27from mindspore.common import dtype as mstype
28from mindspore.dataset.vision import Inter
29from mindspore.nn import Dense, TrainOneStepCell, WithLossCell, ForwardValueAndGrad
30from mindspore.nn.metrics import Accuracy
31from mindspore.nn.optim import Momentum
32from mindspore.ops import operations as P
33from mindspore.ops import functional as F
34from mindspore.train import Model
35from mindspore.train.callback import LossMonitor
36from mindspore.common.initializer import TruncatedNormal
37
38context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
39
40
41def conv(in_channels, out_channels, kernel_size, stride=1, padding=0):
42    """weight initial for conv layer"""
43    weight = weight_variable()
44    return nn.Conv2d(in_channels, out_channels,
45                     kernel_size=kernel_size, stride=stride, padding=padding,
46                     weight_init=weight, has_bias=False, pad_mode="valid")
47
48
49def fc_with_initialize(input_channels, out_channels):
50    """weight initial for fc layer"""
51    weight = weight_variable()
52    bias = weight_variable()
53    return nn.Dense(input_channels, out_channels, weight, bias)
54
55
56def weight_variable():
57    """weight initial"""
58    return TruncatedNormal(0.02)
59
60
61class LeNet5(nn.Cell):
62    def __init__(self, num_class=10, channel=1):
63        super(LeNet5, self).__init__()
64        self.num_class = num_class
65        self.conv1 = conv(channel, 6, 5)
66        self.conv2 = conv(6, 16, 5)
67        self.fc1 = fc_with_initialize(16 * 5 * 5, 120)
68        self.fc2 = fc_with_initialize(120, 84)
69        self.fc3 = fc_with_initialize(84, self.num_class)
70        self.relu = nn.ReLU()
71        self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)
72        self.flatten = nn.Flatten()
73
74    def construct(self, x):
75        x = self.conv1(x)
76        x = self.relu(x)
77        x = self.max_pool2d(x)
78        x = self.conv2(x)
79        x = self.relu(x)
80        x = self.max_pool2d(x)
81        x = self.flatten(x)
82        x = self.fc1(x)
83        x = self.relu(x)
84        x = self.fc2(x)
85        x = self.relu(x)
86        x = self.fc3(x)
87        return x
88
89
90class LeNet(nn.Cell):
91    def __init__(self):
92        super(LeNet, self).__init__()
93        self.relu = P.ReLU()
94        self.batch_size = 1
95        weight1 = Tensor(np.ones([6, 3, 5, 5]).astype(np.float32) * 0.01)
96        weight2 = Tensor(np.ones([16, 6, 5, 5]).astype(np.float32) * 0.01)
97        self.conv1 = nn.Conv2d(3, 6, (5, 5), weight_init=weight1, stride=1, padding=0, pad_mode='valid')
98        self.conv2 = nn.Conv2d(6, 16, (5, 5), weight_init=weight2, pad_mode='valid', stride=1, padding=0)
99        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, pad_mode="valid")
100
101        self.reshape = P.Reshape()
102        self.reshape1 = P.Reshape()
103
104        self.fc1 = Dense(400, 120)
105        self.fc2 = Dense(120, 84)
106        self.fc3 = Dense(84, 10)
107
108    def construct(self, input_x):
109        output = self.conv1(input_x)
110        output = self.relu(output)
111        output = self.pool(output)
112        output = self.conv2(output)
113        output = self.relu(output)
114        output = self.pool(output)
115        output = self.reshape(output, (self.batch_size, -1))
116        output = self.fc1(output)
117        output = self.fc2(output)
118        output = self.fc3(output)
119        return output
120
121
122def multisteplr(total_steps, gap, base_lr=0.9, gamma=0.1, dtype=mstype.float32):
123    lr = []
124    for step in range(total_steps):
125        lr_ = base_lr * gamma ** (step // gap)
126        lr.append(lr_)
127    return Tensor(np.array(lr), dtype)
128
129
130@pytest.mark.level0
131@pytest.mark.platform_x86_gpu_training
132@pytest.mark.env_onecard
133def test_train_lenet():
134    epoch = 100
135    net = LeNet()
136    momentum = 0.9
137    learning_rate = multisteplr(epoch, 30)
138
139    optimizer = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), learning_rate, momentum)
140    criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
141    net_with_criterion = WithLossCell(net, criterion)
142    train_network = TrainOneStepCell(net_with_criterion, optimizer)  # optimizer
143    train_network.set_train()
144    losses = []
145    for i in range(epoch):
146        data = Tensor(np.ones([net.batch_size, 3, 32, 32]).astype(np.float32) * 0.01)
147        label = Tensor(np.ones([net.batch_size]).astype(np.int32))
148        loss = train_network(data, label).asnumpy()
149        losses.append(loss)
150    assert losses[-1] < 0.01
151
152
153def create_dataset(data_path, batch_size=32, repeat_size=1,
154                   num_parallel_workers=1):
155    """
156    create dataset for train or test
157    """
158    # define dataset
159    mnist_ds = ds.MnistDataset(data_path)
160
161    resize_height, resize_width = 32, 32
162    rescale = 1.0 / 255.0
163    shift = 0.0
164    rescale_nml = 1 / 0.3081
165    shift_nml = -1 * 0.1307 / 0.3081
166
167    # define map operations
168    resize_op = CV.Resize((resize_height, resize_width), interpolation=Inter.LINEAR)  # Bilinear mode
169    rescale_nml_op = CV.Rescale(rescale_nml, shift_nml)
170    rescale_op = CV.Rescale(rescale, shift)
171    hwc2chw_op = CV.HWC2CHW()
172    type_cast_op = C.TypeCast(mstype.int32)
173
174    # apply map operations on images
175    mnist_ds = mnist_ds.map(operations=type_cast_op, input_columns="label", num_parallel_workers=num_parallel_workers)
176    mnist_ds = mnist_ds.map(operations=resize_op, input_columns="image", num_parallel_workers=num_parallel_workers)
177    mnist_ds = mnist_ds.map(operations=rescale_op, input_columns="image", num_parallel_workers=num_parallel_workers)
178    mnist_ds = mnist_ds.map(operations=rescale_nml_op, input_columns="image", num_parallel_workers=num_parallel_workers)
179    mnist_ds = mnist_ds.map(operations=hwc2chw_op, input_columns="image", num_parallel_workers=num_parallel_workers)
180
181    # apply DatasetOps
182    buffer_size = 10000
183    mnist_ds = mnist_ds.shuffle(buffer_size=buffer_size)  # 10000 as in LeNet train script
184    mnist_ds = mnist_ds.batch(batch_size, drop_remainder=True)
185    mnist_ds = mnist_ds.repeat(repeat_size)
186
187    return mnist_ds
188
189
190@pytest.mark.level1
191@pytest.mark.platform_x86_gpu_training
192@pytest.mark.env_onecard
193def test_train_and_eval_lenet():
194    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
195    network = LeNet5(10)
196    net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
197    net_opt = nn.Momentum(network.trainable_params(), 0.01, 0.9)
198    model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()})
199
200    print("============== Starting Training ==============")
201    ds_train = create_dataset(os.path.join('/home/workspace/mindspore_dataset/mnist', "train"), 32, 1)
202    model.train(1, ds_train, callbacks=[LossMonitor()], dataset_sink_mode=True)
203
204    print("============== Starting Testing ==============")
205    ds_eval = create_dataset(os.path.join('/home/workspace/mindspore_dataset/mnist', "test"), 32, 1)
206    acc = model.eval(ds_eval, dataset_sink_mode=True)
207    print("============== {} ==============".format(acc))
208
209
210@pytest.mark.level0
211@pytest.mark.platform_x86_gpu_training
212@pytest.mark.env_onecard
213def test_train_lenet_with_new_interface(num_classes=10, epoch=20, batch_size=32):
214    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
215    network = LeNet5(num_classes)
216    criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
217    net_with_criterion = WithLossCell(network, criterion)
218    net_with_criterion.set_train()
219
220    weights = ParameterTuple(network.trainable_params())
221    optimizer = nn.Momentum(weights, 0.1, 0.9)
222
223    train_network = ForwardValueAndGrad(network=net_with_criterion, weights=weights, get_by_list=True, sens_param=True)
224    losses = []
225    for i in range(0, epoch):
226        data = Tensor(np.ones([batch_size, 1, 32, 32]).astype(np.float32) * 0.01)
227        label = Tensor(np.ones([batch_size]).astype(np.int32))
228        sens = Tensor(np.ones([1]).astype(np.float32))
229        loss, grads = train_network(data, label, sens)
230        grads = F.identity(grads)
231        optimizer(grads)
232        losses.append(loss)
233    assert losses[-1].asnumpy() < 0.01
234    assert losses[-1].asnumpy() > 0.001
235