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