def criterion_parallel_apply(module, inputs, targets, kwargs_tup=None, sync=False): """Data Parallel Criterion""" if kwargs_tup: assert len(inputs) == len(kwargs_tup) else: kwargs_tup = ({}, ) * len(inputs) lock = threading.Lock() results = {} def _worker(i, module, input, target, kwargs, results, is_recording, is_training, lock): try: if is_recording: with autograd.record(is_training): output = module(*(input + target), **kwargs) output.wait_to_read() else: output = module(*(input + target), **kwargs) output.wait_to_read() with lock: results[i] = output except Exception as e: with lock: results[i] = e is_training = bool(autograd.is_training()) is_recording = autograd.is_recording() threads = [ threading.Thread( target=_worker, args=(i, module, input, target, kwargs, results, is_recording, is_training, lock), ) for i, (input, target, kwargs) in enumerate(zip(inputs, targets, kwargs_tup)) ] if sync: for thread in threads: thread.start() for thread in threads: thread.join() outputs = [] for i in range(len(inputs)): output = results[i] if isinstance(output, Exception): raise output outputs.append(output) return tuple(outputs) else: outputs = [module(*(input + target), **kwargs) \ for (input, target, kwargs) in zip(inputs, targets, kwargs_tup)] return tuple(outputs)
def parallel_apply(module, inputs, kwargs_tup=None, sync=False): """Parallel applying model forward""" if kwargs_tup is not None: assert len(inputs) == len(kwargs_tup) else: kwargs_tup = ({}, ) * len(inputs) lock = threading.Lock() results = {} def _worker(i, module, input, kwargs, results, is_recording, is_training, lock): try: if is_recording: with autograd.record(is_training): output = tuple_map(module(*input, **kwargs)) for out in output: out.wait_to_read() else: output = tuple_map(module(*input, **kwargs)) for out in output: out.wait_to_read() with lock: results[i] = output except Exception as e: with lock: results[i] = e is_training = autograd.is_training() is_recording = autograd.is_recording() threads = [ threading.Thread( target=_worker, args=(i, module, input, kwargs, results, is_recording, is_training, lock), ) for i, (input, kwargs) in enumerate(zip(inputs, kwargs_tup)) ] if sync: for thread in threads: thread.start() for thread in threads: thread.join() outputs = [] for i in range(len(inputs)): output = results[i] if isinstance(output, Exception): raise output outputs.append(output) return tuple(outputs) else: outputs = [ tuple_map(module(*input, **kwargs)) for (input, kwargs) in zip(inputs, kwargs_tup) ] return tuple(outputs)
def hybrid_forward(self, F, x): flag = autograd.is_recording() e = self.e(x) result, start, end = [], 0, 0 for i, size in enumerate(self.d): start, end = end, end + size t1, t2 = start, end if flag and size > 5: t1, t2 = randomRange(start, end) sliced = F.slice_axis(e, 1, t1, t2) result.append(F.mean(sliced, 1, True)) return F.concat(*result, dim=1)
def hybrid_forward(self, F, x): """Hybrid forward""" features = self.features(x) cls_preds = [ F.flatten(F.transpose(cp(feat), (0, 2, 3, 1))) for feat, cp in zip(features, self.class_predictors) ] box_preds = [ F.flatten(F.transpose(bp(feat), (0, 2, 3, 1))) for feat, bp in zip(features, self.box_predictors) ] anchors = [ F.reshape(ag(feat), shape=(1, -1)) for feat, ag in zip(features, self.anchor_generators) ] cls_preds = F.concat(*cls_preds, dim=1).reshape( (0, -1, self.num_classes)) box_preds = F.concat(*box_preds, dim=1).reshape((0, -1, 4)) anchors = F.concat(*anchors, dim=1).reshape((1, -1, 4)) if autograd.is_recording(): return [cls_preds, box_preds, anchors] bboxes = self.bbox_decoder(box_preds, anchors) cls_ids, scores = self.cls_decoder(F.softmax(cls_preds, axis=-1)) results = [] for i in range(self.num_classes - 1): cls_id = cls_ids.slice_axis(axis=-1, begin=i, end=i + 1) score = scores.slice_axis(axis=-1, begin=i, end=i + 1) # per class results per_result = F.concat(*[cls_id, score, bboxes], dim=-1) if self.nms_thresh > 0 and self.nms_thresh < 1: per_result = F.contrib.box_nms(per_result, overlap_thresh=self.nms_thresh, topk=self.nms_topk, id_index=0, score_index=1, coord_start=2) results.append(per_result) result = F.concat(*results, dim=1) ids = F.slice_axis(result, axis=2, begin=0, end=1) scores = F.slice_axis(result, axis=2, begin=1, end=2) bboxes = F.slice_axis(result, axis=2, begin=2, end=6) return ids, scores, bboxes
def hybrid_forward(self, F, x, *args): if len(args) != 0 and not autograd.is_training(): raise TypeError('CenterNet inference only need one input data.') feats = self.backbone(x) feat = self.neck(feats) # head convolutions tl_cnv = self.tl_cnvs(feat) br_cnv = self.br_cnvs(feat) ct_cnv = self.ct_cnvs(feat) # heatmap tl_heat = self.tl_heats(tl_cnv) br_heat = self.br_heats(br_cnv) ct_heat = self.ct_heats(ct_cnv) # tags tl_tag = self.tl_tags(tl_cnv) br_tag = self.br_tags(br_cnv) # regs tl_reg = self.tl_regs(tl_cnv) br_reg = self.br_regs(br_cnv) ct_reg = self.ct_regs(ct_cnv) if autograd.is_recording(): # compose & gather the tag & ind feats tl_tag = self.transpose_gather_feats(tl_tag, args[0]) br_tag = self.transpose_gather_feats(br_tag, args[1]) tl_reg = self.transpose_gather_feats(tl_reg, args[0]) br_reg = self.transpose_gather_feats(br_reg, args[1]) ct_reg = self.transpose_gather_feats(ct_reg, args[2]) outs = [ tl_heat, br_heat, ct_heat, tl_tag, br_tag, tl_reg, br_reg, ct_reg ] else: outs = [ tl_heat, br_heat, ct_heat, tl_tag, br_tag, tl_reg, br_reg, ct_reg ] return outs
def hybrid_forward(self, F, x, *args): """YOLOV3 network hybrid forward. Parameters ---------- F : mxnet.nd or mxnet.sym `F` is mxnet.sym if hybridized or mxnet.nd if not. x : mxnet.nd.NDArray Input data. *args : optional, mxnet.nd.NDArray During training, extra inputs are required: (gt_boxes, obj_t, centers_t, scales_t, weights_t, clas_t) These are generated by YOLOV3PrefetchTargetGenerator in dataloader transform function. Returns ------- (tuple of) mxnet.nd.NDArray During inference, return detections in shape (B, N, 6) with format (cid, score, xmin, ymin, xmax, ymax) During training, return losses only: (obj_loss, center_loss, scale_loss, cls_loss). """ all_box_centers = [] all_box_scales = [] all_objectness = [] all_class_pred = [] all_anchors = [] all_offsets = [] all_feat_maps = [] all_detections = [] routes = [] for stage, block, output in zip(self.stages, self.yolo_blocks, self.yolo_outputs): x = stage(x) routes.append(x) # the YOLO output layers are used in reverse order, i.e., from very deep layers to shallow for i, block, output in zip(range(len(routes)), self.yolo_blocks, self.yolo_outputs): x, tip = block(x) if autograd.is_training(): dets, box_centers, box_scales, objness, class_pred, anchors, offsets = output( tip) all_box_centers.append(box_centers.reshape((0, -3, -1))) all_box_scales.append(box_scales.reshape((0, -3, -1))) all_objectness.append(objness.reshape((0, -3, -1))) all_class_pred.append(class_pred.reshape((0, -3, -1))) all_anchors.append(anchors) all_offsets.append(offsets) # here we use fake featmap to reduce memory consuption, only shape[2, 3] is used fake_featmap = F.zeros_like( tip.slice_axis(axis=0, begin=0, end=1).slice_axis(axis=1, begin=0, end=1)) all_feat_maps.append(fake_featmap) else: dets = output(tip) all_detections.append(dets) if i >= len(routes) - 1: break # add transition layers x = self.transitions[i](x) # upsample feature map reverse to shallow layers upsample = _upsample(x, stride=2) route_now = routes[::-1][i + 1] x = F.concat(F.slice_like(upsample, route_now * 0, axes=(2, 3)), route_now, dim=1) if autograd.is_training(): # during training, the network behaves differently since we don't need detection results if autograd.is_recording(): # generate losses and return them directly box_preds = F.concat(*all_detections, dim=1) all_preds = [ F.concat(*p, dim=1) for p in [ all_objectness, all_box_centers, all_box_scales, all_class_pred ] ] all_targets = self._target_generator(box_preds, *args) return self._loss(*(all_preds + all_targets)) # return raw predictions, this is only used in DataLoader transform function. return (F.concat(*all_detections, dim=1), all_anchors, all_offsets, all_feat_maps, F.concat(*all_box_centers, dim=1), F.concat(*all_box_scales, dim=1), F.concat(*all_objectness, dim=1), F.concat(*all_class_pred, dim=1)) # concat all detection results from different stages result = F.concat(*all_detections, dim=1) # apply nms per class if 0 < self.nms_thresh < 1: result = F.contrib.box_nms(result, overlap_thresh=self.nms_thresh, valid_thresh=0.01, topk=self.nms_topk, id_index=0, score_index=1, coord_start=2, force_suppress=False) if self.post_nms > 0: result = result.slice_axis(axis=1, begin=0, end=self.post_nms) ids = result.slice_axis(axis=-1, begin=0, end=1) scores = result.slice_axis(axis=-1, begin=1, end=2) bboxes = result.slice_axis(axis=-1, begin=2, end=6) return ids, scores, bboxes
def hybrid_forward(self, F, x, *args): """FYOLO network hybrid forward. Parameters ---------- F : mxnet.nd or mxnet.sym `F` is mxnet.sym if hybridized or mxnet.nd if not. x : mxnet.nd.NDArray Input data. *args : optional, mxnet.nd.NDArray During training, extra inputs are required: (gt_boxes, obj_t, centers_t, scales_t, weights_t, clas_t) These are generated by YOLOV3PrefetchTargetGenerator in dataloader transform function. Returns ------- (tuple of) mxnet.nd.NDArray During inference, return detections in shape (B, N, 6) with format (cid, score, xmin, ymin, xmax, ymax) During training, return losses only: (obj_loss, center_loss, scale_loss, cls_loss). """ x = self.feature(x) x = self.yolo_block(x) all_box_centers, all_box_scales, all_objectness, all_class_pred, all_anchors, \ all_offsets, all_feat_maps, all_detections = [], [], [], [], [], [], [] if autograd.is_training(): dets, box_centers, box_scales, objness, class_pred, anchors, offsets = self.yolo_output(x) all_box_centers.append(box_centers.reshape((0, -3, -1))) all_box_scales.append(box_scales.reshape((0, -3, -1))) all_objectness.append(objness.reshape((0, -3, -1))) all_class_pred.append(class_pred.reshape((0, -3, -1))) all_anchors.append(anchors) all_offsets.append(offsets) # here we use fake featmap to reduce memory consuption, only shape[2, 3] is used fake_featmap = F.zeros_like(x.slice_axis(axis=0, begin=0, end=1).slice_axis(axis=1, begin=0, end=1)) all_feat_maps.append(fake_featmap) else: dets = self.yolo_output(x) all_detections.append(dets) if autograd.is_training(): # during training, the network behaves differently since we don't need detection results if autograd.is_recording(): # generate losses and return them directly box_preds = F.concat(*all_detections, dim=1) all_preds = [F.concat(*p, dim=1) for p in [ all_objectness, all_box_centers, all_box_scales, all_class_pred]] all_targets = self._target_generator(box_preds, *args) return self._loss(*(all_preds + all_targets)) # return raw predictions, this is only used in DataLoader transform function. return (F.concat(*all_detections, dim=1), all_anchors, all_offsets, all_feat_maps, F.concat(*all_box_centers, dim=1), F.concat(*all_box_scales, dim=1), F.concat(*all_objectness, dim=1), F.concat(*all_class_pred, dim=1)) # concat all detection results from different stages result = F.concat(*all_detections, dim=1) # apply nms per class if self.nms_thresh > 0 and self.nms_thresh < 1: result = F.contrib.box_nms( result, overlap_thresh=self.nms_thresh, valid_thresh=0.01, topk=self.nms_topk, id_index=0, score_index=1, coord_start=2, force_suppress=False) if self.post_nms > 0: result = result.slice_axis(axis=1, begin=0, end=self.post_nms) ids = result.slice_axis(axis=-1, begin=0, end=1) scores = result.slice_axis(axis=-1, begin=1, end=2) bboxes = result.slice_axis(axis=-1, begin=2, end=None) return ids, scores, bboxes
def hybrid_forward(self, F, x, *args): """YOLOV3 network hybrid forward. Parameters ---------- F : mxnet.nd or mxnet.sym `F` is mxnet.sym if hybridized or mxnet.nd if not. x : mxnet.nd.NDArray Input data. *args : optional, mxnet.nd.NDArray During training, extra inputs are required: (gt_boxes, obj_t, centers_t, scales_t, weights_t, clas_t) These are generated by YOLOV3PrefetchTargetGenerator in dataloader transform function. Returns ------- (tuple of) mxnet.nd.NDArray During inference, return detections in shape (B, N, 6) with format (cid, score, xmin, ymin, xmax, ymax) During training, return losses only: (obj_loss, center_loss, scale_loss, cls_loss). """ all_box_centers = [] all_box_scales = [] all_objectness = [] all_class_pred = [] all_anchors = [] all_offsets = [] all_feat_maps = [] all_detections = [] routes = [] for stage, block, output in zip(self.stages, self.yolo_blocks, self.yolo_outputs): x = stage(x) routes.append(x) # the YOLO output layers are used in reverse order, i.e., from very deep layers to shallow for i, block, output in zip(range(len(routes)), self.yolo_blocks, self.yolo_outputs): x, tip = block(x) if autograd.is_training(): dets, box_centers, box_scales, objness, class_pred, anchors, offsets = output(tip) all_box_centers.append(box_centers.reshape((0, -3, -1))) all_box_scales.append(box_scales.reshape((0, -3, -1))) all_objectness.append(objness.reshape((0, -3, -1))) all_class_pred.append(class_pred.reshape((0, -3, -1))) all_anchors.append(anchors) all_offsets.append(offsets) # here we use fake featmap to reduce memory consuption, only shape[2, 3] is used fake_featmap = F.zeros_like(tip.slice_axis( axis=0, begin=0, end=1).slice_axis(axis=1, begin=0, end=1)) all_feat_maps.append(fake_featmap) else: dets = output(tip) all_detections.append(dets) if i >= len(routes) - 1: break # add transition layers x = self.transitions[i](x) # upsample feature map reverse to shallow layers upsample = _upsample(x, stride=2) route_now = routes[::-1][i + 1] x = F.concat(F.slice_like(upsample, route_now * 0, axes=(2, 3)), route_now, dim=1) if autograd.is_training(): # during training, the network behaves differently since we don't need detection results if autograd.is_recording(): # generate losses and return them directly box_preds = F.concat(*all_detections, dim=1) all_preds = [F.concat(*p, dim=1) for p in [ all_objectness, all_box_centers, all_box_scales, all_class_pred]] all_targets = self._target_generator(box_preds, *args) return self._loss(*(all_preds + all_targets)) # return raw predictions, this is only used in DataLoader transform function. return (F.concat(*all_detections, dim=1), all_anchors, all_offsets, all_feat_maps, F.concat(*all_box_centers, dim=1), F.concat(*all_box_scales, dim=1), F.concat(*all_objectness, dim=1), F.concat(*all_class_pred, dim=1)) # concat all detection results from different stages result = F.concat(*all_detections, dim=1) # apply nms per class if self.nms_thresh > 0 and self.nms_thresh < 1: result = F.contrib.box_nms( result, overlap_thresh=self.nms_thresh, valid_thresh=0.01, topk=self.nms_topk, id_index=0, score_index=1, coord_start=2, force_suppress=False) if self.post_nms > 0: result = result.slice_axis(axis=1, begin=0, end=self.post_nms) ids = result.slice_axis(axis=-1, begin=0, end=1) scores = result.slice_axis(axis=-1, begin=1, end=2) bboxes = result.slice_axis(axis=-1, begin=2, end=None) return ids, scores, bboxes
def forward_propagation(self, F, x): ''' Hybrid forward with Gluon API. Args: F: `mxnet.ndarray` or `mxnet.symbol`. x: `mxnet.ndarray` of observed data points. Returns: `mxnet.ndarray` or `mxnet.symbol` of inferenced feature points. ''' noised_x = self.noiseable_data.noise(x, F=F) for i in range(len(self.encoder.hidden_units_list)): x = self.encoder.hidden_units_list[i](x) noised_x = self.encoder.hidden_units_list[i](noised_x) if self.encoder.hidden_activation_list[i] == "identity_adjusted": x = x / F.sum(F.ones_like(x)) noised_x = noised_x / F.sum(F.ones_like(noised_x)) elif self.encoder.hidden_activation_list[i] != "identity": x = F.Activation(x, self.encoder.hidden_activation_list[i]) noised_x = F.Activation(noised_x, self.encoder.hidden_activation_list[i]) if self.encoder.hidden_dropout_rate_list[i] is not None: x = self.encoder.hidden_dropout_rate_list[i](x) noised_x = self.encoder.hidden_dropout_rate_list[i](noised_x) if self.encoder.hidden_batch_norm_list[i] is not None: x = self.encoder.hidden_batch_norm_list[i](x) noised_x = self.encoder.hidden_batch_norm_list[i](noised_x) noised_x = self.noiseable_data.noise(noised_x, F=F) alpha_arr, sigma_arr, mu_arr = self.forward_ladder_net(F, x, noised_x) x = F.broadcast_add(x, (alpha_arr * self.alpha)) x = F.broadcast_add(x, (sigma_arr * self.sigma)) x = F.broadcast_add(x, (mu_arr * self.mu)) if self.output_nn is None: self.feature_points_arr = x else: inner_x = self.output_nn.forward_propagation(F, x) self.feature_points_arr = inner_x if self.recoding_ld_loss is True: if autograd.is_recording(): self.alpha_encoder_loss_list.append(F.mean(F.square(alpha_arr)).asnumpy()) self.sigma_encoder_loss_list.append(F.mean(F.square(sigma_arr)).asnumpy()) self.mu_encoder_loss_list.append(F.mean(F.square(mu_arr)).asnumpy()) else: self.alpha_encoder_test_loss_list.append(F.mean(F.square(alpha_arr)).asnumpy()) self.sigma_encoder_test_loss_list.append(F.mean(F.square(sigma_arr)).asnumpy()) self.mu_encoder_test_loss_list.append(F.mean(F.square(mu_arr)).asnumpy()) noised_x = self.noiseable_data.noise(x, F=F) for i in range(len(self.decoder.hidden_units_list)): x = self.decoder.hidden_units_list[i](x) noised_x = self.decoder.hidden_units_list[i](noised_x) if self.decoder.hidden_activation_list[i] == "identity_adjusted": x = x / F.sum(F.ones_like(x)) noised_x = noised_x / F.sum(F.ones_like(noised_x)) elif self.decoder.hidden_activation_list[i] != "identity": x = F.Activation(x, self.decoder.hidden_activation_list[i]) noised_x = F.Activation(noised_x, self.decoder.hidden_activation_list[i]) if self.decoder.hidden_dropout_rate_list[i] is not None: x = self.decoder.hidden_dropout_rate_list[i](x) noised_x = self.decoder.hidden_dropout_rate_list[i](noised_x) if self.decoder.hidden_batch_norm_list[i] is not None: x = self.decoder.hidden_batch_norm_list[i](x) noised_x = self.decoder.hidden_batch_norm_list[i](noised_x) noised_x = self.noiseable_data.noise(noised_x, F=F) alpha_arr, sigma_arr, mu_arr = self.forward_ladder_net(F, x, noised_x) x = F.broadcast_add(x, (alpha_arr * self.alpha)) x = F.broadcast_add(x, (sigma_arr * self.sigma)) x = F.broadcast_add(x, (mu_arr * self.mu)) if self.decoder.output_nn is not None: x = self.decoder.output_nn.forward_propagation(F, x) if self.recoding_ld_loss is True: if autograd.is_recording(): self.alpha_decoder_loss_list.append(F.mean(F.square(alpha_arr)).asnumpy()) self.sigma_decoder_loss_list.append(F.mean(F.square(sigma_arr)).asnumpy()) self.mu_decoder_loss_list.append(F.mean(F.square(mu_arr)).asnumpy()) else: self.alpha_decoder_test_loss_list.append(F.mean(F.square(alpha_arr)).asnumpy()) self.sigma_decoder_test_loss_list.append(F.mean(F.square(sigma_arr)).asnumpy()) self.mu_decoder_test_loss_list.append(F.mean(F.square(mu_arr)).asnumpy()) return x
def inference_g(self, observed_arr): ''' Inference with generator. Args: observed_arr: `mxnet.ndarray` of observed data points. Returns: Tuple data. - re-parametric data. - encoded data points. - re-encoded data points. ''' generated_arr, encoded_arr, re_encoded_arr = super().inference_g(observed_arr) if autograd.is_recording(): limit = self.long_term_seq_len seq_len = self.noise_sampler.seq_len self.noise_sampler.seq_len = limit long_term_observed_arr = self.noise_sampler.draw() observed_mean_arr = nd.expand_dims(nd.mean(long_term_observed_arr, axis=1), axis=1) sum_arr = None for seq in range(2, long_term_observed_arr.shape[1]): add_arr = nd.sum(long_term_observed_arr[:, :seq] - observed_mean_arr, axis=1) if sum_arr is None: sum_arr = nd.expand_dims(add_arr, axis=0) else: sum_arr = nd.concat( sum_arr, nd.expand_dims(add_arr, axis=0), dim=0 ) max_arr = nd.max(sum_arr, axis=0) min_arr = nd.min(sum_arr, axis=0) diff_arr = long_term_observed_arr - observed_mean_arr std_arr = nd.power(nd.mean(nd.square(diff_arr), axis=1), 1/2) R_S_arr = (max_arr - min_arr) / std_arr len_arr = nd.ones_like(R_S_arr, ctx=R_S_arr.context) * np.log(long_term_observed_arr.shape[1] / 2) observed_H_arr = nd.log(R_S_arr) / len_arr self.noise_sampler.seq_len = seq_len g_min_arr = nd.expand_dims(generated_arr.min(axis=1), axis=1) g_max_arr = nd.expand_dims(generated_arr.max(axis=1), axis=1) o_min_arr = nd.expand_dims(observed_arr.min(axis=1), axis=1) o_max_arr = nd.expand_dims(observed_arr.max(axis=1), axis=1) _observed_arr = generated_arr long_term_generated_arr = None for i in range(limit): generated_arr, _, _ = super().inference_g(_observed_arr) g_min_arr = nd.expand_dims(generated_arr.min(axis=1), axis=1) g_max_arr = nd.expand_dims(generated_arr.max(axis=1), axis=1) o_min_arr = nd.expand_dims(_observed_arr.min(axis=1), axis=1) o_max_arr = nd.expand_dims(_observed_arr.max(axis=1), axis=1) generated_arr = (generated_arr - g_min_arr) / (g_max_arr - g_min_arr) generated_arr = (o_max_arr - o_min_arr) * generated_arr generated_arr = o_min_arr + generated_arr if self.condition_sampler is not None: self.condition_sampler.output_shape = generated_arr.shape noise_arr = self.condition_sampler.generate() generated_arr += noise_arr if long_term_generated_arr is None: long_term_generated_arr = generated_arr else: long_term_generated_arr = nd.concat( long_term_generated_arr, generated_arr, dim=1 ) _observed_arr = generated_arr generated_mean_arr = nd.expand_dims(nd.mean(long_term_generated_arr, axis=1), axis=1) sum_arr = None for seq in range(2, long_term_generated_arr.shape[1]): add_arr = nd.sum(long_term_generated_arr[:, :seq] - generated_mean_arr, axis=1) if sum_arr is None: sum_arr = nd.expand_dims(add_arr, axis=0) else: sum_arr = nd.concat( sum_arr, nd.expand_dims(add_arr, axis=0), dim=0 ) max_arr = nd.max(sum_arr, axis=0) min_arr = nd.min(sum_arr, axis=0) diff_arr = long_term_generated_arr - generated_mean_arr std_arr = nd.power(nd.mean(nd.square(diff_arr), axis=1), 1/2) R_S_arr = (max_arr - min_arr) / std_arr len_arr = nd.ones_like(R_S_arr, ctx=R_S_arr.context) * np.log(long_term_generated_arr.shape[1] / 2) generated_H_arr = nd.log(R_S_arr) / len_arr multi_fractal_loss = nd.abs(generated_H_arr - observed_H_arr) multi_fractal_loss = nd.mean(multi_fractal_loss, axis=0, exclude=True) multi_fractal_loss = nd.expand_dims(multi_fractal_loss, axis=-1) multi_fractal_loss = nd.expand_dims(multi_fractal_loss, axis=-1) generated_arr = generated_arr + multi_fractal_loss return generated_arr, encoded_arr, re_encoded_arr