• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2020-2021 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"""test_image_summary"""
16import logging
17import os
18import numpy as np
19
20import mindspore.nn as nn
21from mindspore import Model, context
22from mindspore import Tensor
23from mindspore.nn.optim import Momentum
24from mindspore.train.summary.summary_record import SummaryRecord, _cache_summary_tensor_data
25from mindspore.train.callback import Callback
26from tests.security_utils import security_off_wrap
27from .....dataset_mock import MindData
28
29CUR_DIR = os.getcwd()
30SUMMARY_DIR = CUR_DIR + "/test_temp_summary_event_file/"
31
32log = logging.getLogger("test")
33log.setLevel(level=logging.ERROR)
34
35
36def make_image_tensor(shape, dtype=float):
37    """ make_image_tensor """
38    number = np.prod(shape)
39    x = (np.arange(number, dtype=dtype)).reshape(shape)
40    return x
41
42
43def get_test_data(step):
44    """ get_test_data """
45    test_data_list = []
46    tag1 = "x1[:Image]"
47    tag2 = "x2[:Image]"
48    np1 = make_image_tensor([2, 3, 8, 8])
49    np2 = make_image_tensor([step, 3, 8, 8])
50
51    dict1 = {}
52    dict1["name"] = tag1
53    dict1["data"] = Tensor(np1)
54
55    dict2 = {}
56    dict2["name"] = tag2
57    dict2["data"] = Tensor(np2)
58
59    test_data_list.append(dict1)
60    test_data_list.append(dict2)
61
62    return test_data_list
63
64
65# Test: call method on parse graph code
66@security_off_wrap
67def test_image_summary_sample():
68    """ test_image_summary_sample """
69    with SummaryRecord(SUMMARY_DIR, file_suffix="_MS_IMAGE") as test_writer:
70
71        for i in range(1, 5):
72            test_data = get_test_data(i)
73            _cache_summary_tensor_data(test_data)
74            test_writer.record(i)
75            test_writer.flush()
76
77
78class Net(nn.Cell):
79    """ Net definition """
80
81    def __init__(self):
82        super(Net, self).__init__()
83        self.conv = nn.Conv2d(3, 64, 3, has_bias=False, weight_init='normal',
84                              pad_mode='valid')
85        self.bn = nn.BatchNorm2d(64)
86        self.relu = nn.ReLU()
87        self.flatten = nn.Flatten()
88        self.fc = nn.Dense(64 * 222 * 222, 3)  # padding=0
89
90    def construct(self, x):
91        x = self.conv(x)
92        x = self.bn(x)
93        x = self.relu(x)
94        x = self.flatten(x)
95        out = self.fc(x)
96        return out
97
98
99class LossNet(nn.Cell):
100    """ LossNet definition """
101
102    def __init__(self):
103        super(LossNet, self).__init__()
104        self.conv = nn.Conv2d(3, 64, 3, has_bias=False, weight_init='normal',
105                              pad_mode='valid')
106        self.bn = nn.BatchNorm2d(64)
107        self.relu = nn.ReLU()
108        self.flatten = nn.Flatten()
109        self.fc = nn.Dense(64 * 222 * 222, 3)  # padding=0
110        self.loss = nn.SoftmaxCrossEntropyWithLogits()
111
112    def construct(self, x, y):
113        x = self.conv(x)
114        x = self.bn(x)
115        x = self.relu(x)
116        x = self.flatten(x)
117        x = self.fc(x)
118        out = self.loss(x, y)
119        return out
120
121
122def get_model():
123    """ get_model """
124    net = Net()
125    loss = nn.SoftmaxCrossEntropyWithLogits()
126    optim = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
127    context.set_context(mode=context.GRAPH_MODE)
128    model = Model(net, loss_fn=loss, optimizer=optim, metrics=None)
129    return model
130
131
132def get_dataset():
133    """ get_dataset """
134    dataset_types = (np.float32, np.float32)
135    dataset_shapes = ((2, 3, 224, 224), (2, 3))
136
137    dataset = MindData(size=2, batch_size=2,
138                       np_types=dataset_types,
139                       output_shapes=dataset_shapes,
140                       input_indexs=(0, 1))
141    return dataset
142
143
144class ImageSummaryCallback(Callback):
145    """Image summary callback."""
146
147    def __init__(self, summary_record):
148        self._summary_record = summary_record
149
150    def __enter__(self):
151        return self
152
153    def __exit__(self, *err):
154        self._summary_record.close()
155
156    def record(self, step, train_network=None):
157        """record data."""
158        self._summary_record.record(step, train_network)
159        self._summary_record.flush()
160
161
162@security_off_wrap
163def test_image_summary_train():
164    """ test_image_summary_train """
165    dataset = get_dataset()
166    with SummaryRecord(SUMMARY_DIR, file_suffix="_MS_IMAGE") as test_writer:
167        model = get_model()
168        callback = ImageSummaryCallback(test_writer)
169        model.train(2, dataset, callbacks=[callback])
170
171
172@security_off_wrap
173def test_image_summary_data():
174    """ test_image_summary_data """
175    dataset = get_dataset()
176
177    test_data_list = []
178    i = 1
179    for next_element in dataset:
180        tag = "image_" + str(i) + "[:Image]"
181        dct = {}
182        dct["name"] = tag
183        dct["data"] = Tensor(next_element[0])
184        test_data_list.append(dct)
185        i += 1
186
187    with SummaryRecord(SUMMARY_DIR, file_suffix="_MS_IMAGE") as test_writer:
188        _cache_summary_tensor_data(test_data_list)
189        test_writer.record(1)
190