1# Copyright 2021-2022 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"""Offload Support. 16""" 17import json 18import numpy as np 19 20import mindspore.common.dtype as mstype 21from mindspore.common.tensor import Tensor 22import mindspore.ops as ops 23import mindspore.nn as nn 24import mindspore.ops.composite as C 25from mindspore.ops import operations as P 26from mindspore.ops.primitive import _primexpr 27from mindspore import log as logger 28 29 30def check_add_offload_sink_mode(dataset, dataset_helper, network): 31 """ 32 Check if any map operations were removed to be offloaded and apply the transforms if so. 33 """ 34 if hasattr(dataset, '__no_send__'): 35 # Dataset was not sent to device. Skip adding offload. 36 return network 37 offload_model = dataset.__transfer_dataset__.get_offload_model() 38 # See if the offload pass identified any operations to be offloaded 39 if offload_model.transform_list != []: 40 check_concat_zip_dataset(dataset.__transfer_dataset__) 41 network = ApplyPreTransform(offload_model, network) 42 return network 43 44 45def check_concat_zip_dataset(dataset): 46 """ 47 Check if dataset is concatenated or zipped. 48 """ 49 child_len = len(dataset.children) 50 if child_len == 1: 51 check_concat_zip_dataset(dataset.children[0]) 52 elif child_len > 1: 53 raise RuntimeError("Offload module currently does not support concatenated or zipped datasets.") 54 55 56def get_col_idxs(node_cols, ds_cols): 57 """ 58 Get the index(es) of the input column(s) from the dataset 59 """ 60 col_idxs = [] 61 non_exist_cols = [] 62 # temporary error if multiple node columns 63 if len(node_cols) > 1: 64 raise RuntimeError( 65 "Offload hardware accelerator currently does not support map operations with multiple input columns") 66 for node_col in node_cols: 67 if node_col in ds_cols: 68 col_idxs.append(ds_cols.index(node_col)) 69 else: 70 non_exist_cols.append(node_col) 71 if non_exist_cols: 72 raise RuntimeError( 73 ("The following input column(s) for an offloaded map operation " 74 "do not exist: {}").format(non_exist_cols)) 75 76 return col_idxs 77 78 79def apply_offload_iterators(data, offload_model): 80 """ 81 Apply offload for non sink mode pipeline. 82 """ 83 non_tensor_idxs = [] 84 for i, _ in enumerate(data): 85 if not isinstance(data[i], Tensor): 86 data[i] = Tensor(data[i], dtype=mstype.float32) 87 non_tensor_idxs.append(i) 88 89 data = offload_model(*data) 90 data = list(data) 91 for idx in non_tensor_idxs: 92 data[idx] = data[idx].asnumpy() 93 94 return data 95 96 97@_primexpr 98def check_input_dims(x_shape, required_dim, offload_op_name): 99 """ 100 Check if input has the required number of dimensions for the operation. 101 """ 102 input_dim = len(x_shape) 103 if input_dim is not required_dim: 104 raise ValueError("For %s offload operation, the dimension of input should be %d, but got %d." % 105 (offload_op_name, required_dim, input_dim)) 106 107 108def assign_min_max_params(in_params, center=1): 109 """ 110 Adjust input parameters for ops. 111 """ 112 if isinstance(in_params, (list, tuple)): 113 min_param = in_params[0] 114 max_param = in_params[1] 115 else: 116 min_param = max(0, center - in_params) 117 max_param = center + in_params 118 119 return min_param, max_param 120 121 122class ApplyPreTransform(nn.Cell): 123 """ 124 Concatenates offload model with network. 125 """ 126 127 def __init__(self, transform, model): 128 super(ApplyPreTransform, self).__init__(auto_prefix=False, flags=model.get_flags()) 129 self.transform = transform 130 self.model = model 131 132 def construct(self, *x): 133 data = [] 134 for data_col in x: 135 data.append(data_col) 136 137 data = self.transform(*data) 138 data = self.model(*data) 139 140 return data 141 142 143class IdentityCell(nn.Cell): 144 """ 145 Applies identity transform on given input tensors. 146 """ 147 148 def __init__(self): 149 super(IdentityCell, self).__init__() 150 self.identity = P.Identity() 151 152 def construct(self, x): 153 return self.identity(x) 154 155 156class RandomHorizontalFlip(nn.Cell): 157 """ 158 Applies Random Horizontal Flip transform on given input tensors. 159 """ 160 161 def __init__(self, prob): 162 super(RandomHorizontalFlip, self).__init__() 163 164 self.prob = Tensor(prob, dtype=mstype.float32) 165 166 self.cast = P.Cast() 167 self.shape = P.Shape() 168 self.reshape = P.Reshape() 169 self.h_flip = P.ReverseV2(axis=[2]) 170 self.mul = P.Mul() 171 172 def construct(self, x): 173 174 x = self.cast(x, mstype.float32) 175 x_shape = self.shape(x) 176 check_input_dims(x_shape, 4, 'RandomHorizontalFlip') 177 bs, h, w, c = x_shape 178 179 flip_rand_factor = Tensor(np.random.uniform(size=(bs, 1)), dtype=mstype.float32) 180 flip_rand_factor = self.cast((self.prob > flip_rand_factor), mstype.float32) 181 flip_rand_factor = self.reshape(C.repeat_elements(flip_rand_factor, rep=(h * w * c)), (bs, h, w, c)) 182 183 x_flip = self.h_flip(x) 184 operation = self.mul(x_flip, flip_rand_factor) + self.mul((1 - flip_rand_factor), x) 185 # ops.depend is added to find the RandomHorizontalFlip operator in IR files. 186 depend = ops.depend(operation, "RandomHorizontalFlip") 187 return depend 188 189 190class RandomVerticalFlip(nn.Cell): 191 """ 192 Applies Random Vertical Flip transform on given input tensors. 193 """ 194 195 def __init__(self, prob): 196 super(RandomVerticalFlip, self).__init__() 197 198 self.prob = Tensor(prob, dtype=mstype.float32) 199 200 self.cast = P.Cast() 201 self.shape = P.Shape() 202 self.reshape = P.Reshape() 203 self.h_flip = P.ReverseV2(axis=[1]) 204 self.mul = P.Mul() 205 206 def construct(self, x): 207 208 x = self.cast(x, mstype.float32) 209 x_shape = self.shape(x) 210 check_input_dims(x_shape, 4, 'RandomVerticalFlip') 211 bs, h, w, c = x_shape 212 213 flip_rand_factor = Tensor(np.random.uniform(size=(bs, 1)), dtype=mstype.float32) 214 flip_rand_factor = self.cast((self.prob > flip_rand_factor), mstype.float32) 215 flip_rand_factor = self.reshape(C.repeat_elements(flip_rand_factor, rep=(h * w * c)), (bs, h, w, c)) 216 217 x_flip = self.h_flip(x) 218 operation = self.mul(x_flip, flip_rand_factor) + self.mul((1 - flip_rand_factor), x) 219 # ops.depend is added to find the RandomVerticalFlip operator in IR files. 220 depend = ops.depend(operation, "RandomVerticalFlip") 221 return depend 222 223 224class GenerateRandBatch(nn.Cell): 225 """ 226 Generate batch with random values uniformly selected from [degree_min, degree_max]. 227 """ 228 229 def __init__(self): 230 super(GenerateRandBatch, self).__init__() 231 232 self.ones = P.Ones() 233 self.reshape = P.Reshape() 234 235 def construct(self, degree_min, degree_max, check_rand, shape): 236 237 bs, h, w, c = shape 238 rand_factor = Tensor(np.random.uniform(size=(bs, 1)), dtype=mstype.float32) 239 rand_factor = degree_min + (degree_max - degree_min)*rand_factor 240 degree_factor = degree_min * self.ones((bs, 1), mstype.float32) 241 rand_factor = (check_rand * degree_factor) + (~check_rand * rand_factor) 242 rand_factor = self.reshape(C.repeat_elements(rand_factor, rep=(h * w * c)), (bs, h, w, c)) 243 244 return rand_factor 245 246 247class RandomColorAdjust(nn.Cell): 248 """ 249 Applies Random Color Adjust transform on given input tensors. 250 """ 251 252 def __init__(self, brightness, contrast, saturation, hue): 253 super(RandomColorAdjust, self).__init__() 254 255 self.br_min, self.br_max = assign_min_max_params(brightness) 256 self.cont_min, self.cont_max = assign_min_max_params(contrast) 257 self.sa_min, self.sa_max = assign_min_max_params(saturation) 258 self.hue_min, self.hue_max = assign_min_max_params(hue) 259 260 self.check_rand_br = Tensor(self.br_min == self.br_max) 261 self.check_rand_cont = Tensor(self.cont_min == self.cont_max) 262 self.check_rand_sa = Tensor(self.sa_min == self.sa_max) 263 self.check_rand_hue = Tensor(self.hue_min == self.hue_max) 264 265 self.cast = P.Cast() 266 self.shape = P.Shape() 267 self.reshape = P.Reshape() 268 self.unstack = P.Unstack(axis=-1) 269 self.unstack_dim_1 = P.Unstack(axis=1) 270 self.expand_dims = P.ExpandDims() 271 self.mul = P.Mul() 272 273 self.mean = P.ReduceMean() 274 self.argmaxvalue = P.ArgMaxWithValue(axis=1, keep_dims=False) 275 self.argminvalue = P.ArgMinWithValue(axis=1, keep_dims=False) 276 self.stack = P.Stack(axis=0) 277 self.epsilon = Tensor(np.finfo(np.float32).eps, mstype.float32) 278 self.squeeze = P.Squeeze(axis=0) 279 self.expand_dims = P.ExpandDims() 280 self.gatherd = P.GatherD() 281 self.floor = P.Floor() 282 self.fmod = P.FloorMod() 283 self.abs = P.Abs() 284 self.zeros_like = P.ZerosLike() 285 self.stack_axis_1 = P.Stack(axis=1) 286 self.transpose = P.Transpose() 287 self.ones = P.Ones() 288 self.reshape = P.Reshape() 289 290 self.generate_rand_batch = GenerateRandBatch() 291 292 def construct(self, x): 293 294 x = self.cast(x, mstype.float32) 295 x_shape = self.shape(x) 296 check_input_dims(x_shape, 4, 'RandomColorAdjust') 297 bs, h, w, c = x_shape 298 299 br_rand_factor = self.generate_rand_batch(self.br_min, self.br_max, self.check_rand_br, x_shape) 300 cont_rand_factor = self.generate_rand_batch(self.cont_min, self.cont_max, self.check_rand_cont, x_shape) 301 sat_rand_factor = self.generate_rand_batch(self.sa_min, self.sa_max, self.check_rand_sa, x_shape) 302 303 r_, g_, b_ = self.unstack(x) 304 305 x_gray = 0.2989 * r_ + 0.587 * g_ + 0.114 * b_ 306 x_gray_mean = self.expand_dims(self.mean(x_gray, (1, 2)) + 0.5, -1) 307 x_gray_mean = self.reshape(C.repeat_elements(x_gray_mean, rep=(h * w * c)), (bs, h, w, c)) 308 x_gray = C.repeat_elements(self.expand_dims(x_gray, -1), rep=c, axis=-1) 309 310 # Apply brightness 311 x = self.mul(x, br_rand_factor) 312 x = ops.clip_by_value(x, 0.0, 255.0) 313 314 # Apply contrast 315 x = self.mul(x, cont_rand_factor) + self.mul((1 - cont_rand_factor), x_gray_mean) 316 x = ops.clip_by_value(x, 0.0, 255.0) 317 318 # Apply saturation 319 x = self.mul(x, sat_rand_factor) + self.mul((1 - sat_rand_factor), x_gray) 320 x = ops.clip_by_value(x, 0.0, 255.0) 321 322 # Apply Hue Transform 323 # Convert tensor from rgb to hsv 324 x_swap = self.transpose(x, (0, 3, 1, 2)) / 255.0 325 r, g, b = self.unstack_dim_1(x_swap) 326 _, max_rgb = self.argmaxvalue(x_swap) 327 argmin_rgb, min_rgb = self.argminvalue(x_swap) 328 329 max_min = max_rgb - min_rgb + self.epsilon 330 h1 = (g - r) * 60 / max_min + 60 331 h2 = (b - g) * 60 / max_min + 180 332 h3 = (r - b) * 60 / max_min + 300 333 hue = self.squeeze(self.gatherd(self.stack((h2, h3, h1)), 0, self.expand_dims(argmin_rgb, 0))) 334 s = max_min / (max_rgb + self.epsilon) 335 v = max_rgb 336 337 # Adjust hue 338 hue_rand_factor = Tensor(np.random.uniform(size=(bs, 1)), dtype=mstype.float32) 339 hue_rand_factor = self.hue_min + (self.hue_max - self.hue_min)*hue_rand_factor 340 degree_factor = self.hue_min * self.ones((bs, 1), mstype.float32) 341 hue_rand_factor = (self.check_rand_hue * degree_factor) + (~self.check_rand_hue * hue_rand_factor) 342 hue_rand_factor = self.reshape(C.repeat_elements(hue_rand_factor, rep=(h * w)), (bs, h, w)) 343 hue = hue + (hue_rand_factor * 360.0) 344 345 # Convert tensor from hsv to rgb 346 h_ = (hue - self.floor(hue / 360.0) * 360.0) / 60.0 347 c = s * v 348 x_ = c * (1 - self.abs(self.fmod(h_, 2) - 1)) 349 350 zero_tensor = self.zeros_like(c) 351 y = self.stack((self.stack_axis_1((c, x_, zero_tensor)), 352 self.stack_axis_1((x_, c, zero_tensor)), 353 self.stack_axis_1((zero_tensor, c, x_)), 354 self.stack_axis_1((zero_tensor, x_, c)), 355 self.stack_axis_1((x_, zero_tensor, c)), 356 self.stack_axis_1((c, zero_tensor, x_)), 357 )) 358 359 index = self.expand_dims(self.floor(h_), 1) 360 index = self.cast(self.expand_dims(C.repeat_elements(index, 3, 1), 0), mstype.int32) 361 x_rgb = self.squeeze(self.gatherd(y, 0, index)) 362 x_rgb = x_rgb + C.repeat_elements(self.expand_dims((v - c), 1), 3, 1) 363 364 x_rgb = self.transpose(x, (0, 2, 3, 1)) * 255.0 365 operation = ops.clip_by_value(x, 0.0, 255.0) 366 # ops.depend is added to find the RandomColorAdjust operator in IR files. 367 depend = ops.depend(operation, "RandomColorAdjust") 368 return depend 369 370 371class RandomSharpness(nn.Cell): 372 """ 373 Applies Random Sharpness transform on given input tensors. 374 """ 375 376 def __init__(self, degrees): 377 super(RandomSharpness, self).__init__() 378 379 if isinstance(degrees, (list, tuple)): 380 self.degree_min = degrees[0] 381 self.degree_max = degrees[1] 382 else: 383 self.degree_min = max(0, 1 - degrees) 384 self.degree_max = 1 + degrees 385 386 self.cast = P.Cast() 387 self.shape = P.Shape() 388 self.ones = P.Ones() 389 self.reshape = P.Reshape() 390 self.expand_dims = P.ExpandDims() 391 self.mul = P.Mul() 392 self.transpose = P.Transpose() 393 394 self.check_rand = Tensor(self.degree_min == self.degree_max) 395 396 self.weight = np.array([[1, 1, 1], [1, 5, 1], [1, 1, 1]])/13.0 397 self.weight = np.repeat(self.weight[np.newaxis, :, :], 3, axis=0) 398 self.weight = np.repeat(self.weight[np.newaxis, :, :], 3, axis=0) 399 self.weight = Tensor(self.weight, mstype.float32) 400 401 self.filter = P.Conv2D(out_channel=3, kernel_size=(3, 3), pad_mode='pad', pad=1) 402 403 def construct(self, x): 404 x = self.cast(x, mstype.float32) 405 x_shape = self.shape(x) 406 check_input_dims(x_shape, 4, 'RandomSharpness') 407 bs, h, w, c = x_shape 408 409 degree_rand_factor = Tensor(np.random.uniform(size=(bs, 1)), dtype=mstype.float32) 410 degree_rand_factor = self.degree_min + (self.degree_max - self.degree_min)*degree_rand_factor 411 degree_factor = self.degree_min * self.ones((bs, 1), mstype.float32) 412 degree_rand_factor = (self.check_rand * degree_factor) + (~self.check_rand * degree_rand_factor) 413 degree_rand_factor = self.reshape(C.repeat_elements(degree_rand_factor, rep=(h * w * c)), (bs, h, w, c)) 414 415 x_sharp = self.filter(self.transpose(x, (0, 3, 1, 2)), self.weight) 416 x_sharp = self.transpose(x_sharp, (0, 2, 3, 1)) 417 418 x = self.mul(x, degree_rand_factor) + self.mul((1 - degree_rand_factor), x_sharp) 419 operation = ops.clip_by_value(x, 0.0, 255.0) 420 # ops.depend is added to find the RandomSharpness operator in IR files. 421 depend = ops.depend(operation, "RandomSharpness") 422 return depend 423 424 425class Rescale(nn.Cell): 426 """ 427 Applies Rescale transform on given input tensors. 428 """ 429 430 def __init__(self, rescale, shift): 431 super(Rescale, self).__init__() 432 433 self.rescale = Tensor(rescale, dtype=mstype.float32) 434 self.shift = Tensor(shift, dtype=mstype.float32) 435 436 self.cast = P.Cast() 437 self.mul = P.Mul() 438 439 def construct(self, x): 440 441 x = self.cast(x, mstype.float32) 442 operation = x * self.rescale + self.shift 443 # ops.depend is added to find the Rescale operator in IR files. 444 depend = ops.depend(operation, "Rescale") 445 return depend 446 447 448class HwcToChw(nn.Cell): 449 """ 450 Applies Channel Swap transform on given input tensors. 451 """ 452 453 def __init__(self): 454 super(HwcToChw, self).__init__() 455 self.trans = P.Transpose() 456 self.shape = P.Shape() 457 458 def construct(self, x): 459 x_shape = self.shape(x) 460 check_input_dims(x_shape, 4, 'HwcToChw') 461 operation = self.trans(x, (0, 3, 1, 2)) 462 # ops.depend is added to find the HwcToChw operator in IR files. 463 depend = ops.depend(operation, "HwcToChw") 464 return depend 465 466 467class Normalize(nn.Cell): 468 """ 469 Applies Normalize transform on given input tensors. 470 """ 471 472 def __init__(self, mean, std, is_hwc=True): 473 super(Normalize, self).__init__() 474 if is_hwc is False: 475 mean = np.expand_dims(np.array(mean), (1, 2)) 476 std = np.expand_dims(np.array(std), (1, 2)) 477 self.mean = Tensor(mean, mstype.float32) 478 self.std = Tensor(std, mstype.float32) 479 self.sub = P.Sub() 480 self.div = P.Div() 481 self.cast = P.Cast() 482 483 def construct(self, x): 484 x = self.cast(x, mstype.float32) 485 x = self.sub(x, self.mean) 486 operation = self.div(x, self.std) 487 # ops.depend is added to find the Normalize operator in IR files. 488 depend = ops.depend(operation, "Normalize") 489 return depend 490 491 492class TypeCast(nn.Cell): 493 """ 494 Applies TypeCast transform on given input tensors. 495 """ 496 497 def __init__(self, data_type_str): 498 super(TypeCast, self).__init__() 499 500 self.cast = P.Cast() 501 self.data_type = mstype.typing.str_to_type(data_type_str) 502 503 def construct(self, x): 504 operation = self.cast(x, self.data_type) 505 # ops.depend is added to find the TypeCast operator in IR files. 506 depend = ops.depend(operation, "TypeCast") 507 return depend 508 509 510class OffloadModel(): 511 def __init__(self, func, args_names=None): 512 self.func = func 513 self.args_names = args_names 514 515 516# Dictionary connecting operation name to model 517op_to_model = { 518 "HWC2CHW": OffloadModel(HwcToChw), 519 "HwcToChw": OffloadModel(HwcToChw), 520 "Normalize": OffloadModel(Normalize, ["mean", "std", "is_hwc"]), 521 "RandomColorAdjust": OffloadModel(RandomColorAdjust, ["brightness", "contrast", "saturation", "hue"]), 522 "RandomHorizontalFlip": OffloadModel(RandomHorizontalFlip, ["prob"]), 523 "RandomSharpness": OffloadModel(RandomSharpness, ["degrees"]), 524 "RandomVerticalFlip": OffloadModel(RandomVerticalFlip, ["prob"]), 525 "Rescale": OffloadModel(Rescale, ["rescale", "shift"]), 526 "TypeCast": OffloadModel(TypeCast, ["data_type"]) 527} 528 529 530class GetModelFromJson2Col(nn.Cell): 531 """ 532 Generates offload ME model from offload JSON file for a single map op. 533 """ 534 535 def __init__(self, json_offload, col_idxs): 536 super(GetModelFromJson2Col, self).__init__() 537 self.col_idxs = col_idxs 538 self.me_ops = [] 539 self.input_cols = [] 540 541 # Check if input_culmns attr is empty in Map op. 542 if not self.col_idxs: 543 self.col_idxs = [0] 544 logger.warning( 545 "input_columns attr in map op is not defined, " 546 "so offload op will be applied to first column of dataset.") 547 548 if json_offload is not None: 549 offload_ops = json_offload["operations"] 550 for op in offload_ops: 551 name = op["tensor_op_name"] 552 args = op["tensor_op_params"] 553 op_model = op_to_model[name] 554 op_model_inputs = [] 555 if op_model.args_names is not None: 556 for arg_key in op_model.args_names: 557 op_model_inputs.append(args[arg_key]) 558 559 self.me_ops.append(op_model.func(*op_model_inputs)) 560 else: 561 raise RuntimeError("Offload hardware accelarator cannot be applied for this pipeline.") 562 563 self.cell = nn.SequentialCell(self.me_ops) 564 565 def construct(self, x): 566 # apply single column 567 col_idx = self.col_idxs[0] 568 x[col_idx] = self.cell(x[col_idx]) 569 570 return x 571 572 573class GetOffloadModel(nn.Cell): 574 """ 575 Generates offload ME model. 576 """ 577 578 def __init__(self, dataset_consumer, ds_cols): 579 super(GetOffloadModel, self).__init__() 580 self.transform_list = [] 581 json_offload = json.loads(dataset_consumer.GetOffload()) 582 if json_offload is not None: 583 for node in json_offload: 584 if node["op_type"] == 'Map': 585 ds_col_idxs = get_col_idxs(node["input_columns"], ds_cols) 586 self.transform_list.append(GetModelFromJson2Col(node, ds_col_idxs)) 587 self.transform_list.reverse() 588 589 def construct(self, *x): 590 data = [] 591 for d in x: 592 data.append(d) 593 594 for transform in self.transform_list: 595 data = transform(data) 596 return data 597