| import logging |
| import os |
| import mxnet as mx |
| import numpy as np |
| import scipy.stats |
| import pickle |
| from ..utils import * |
|
|
| class IdentityOp(mx.operator.CustomOp): |
| def __init__(self, logging_prefix="identity", input_debug=False, grad_debug=False): |
| super(IdentityOp, self).__init__() |
| self.logging_prefix=logging_prefix |
| self.input_debug = input_debug |
| self.grad_debug = grad_debug |
|
|
| def forward(self, is_train, req, in_data, out_data, aux): |
| if(self.input_debug): |
| logging.info("%s: in_norm=%f, in_max=%f, in_mean=%f, in_min=%f, in_shape=%s" |
| %(self.logging_prefix, np.linalg.norm(in_data[0].asnumpy()), in_data[0].asnumpy().max(), np.abs(in_data[0].asnumpy()).mean(), in_data[0].asnumpy().min(), str(in_data[0].shape))) |
| self.assign(out_data[0], req[0], in_data[0]) |
|
|
| def backward(self, req, out_grad, in_data, out_data, in_grad, aux): |
| self.assign(in_grad[0], req[0], out_grad[0]) |
| if (self.grad_debug): |
| logging.info("%s: grad_norm=%f, grad_shape=%s" |
| % (self.logging_prefix, np.linalg.norm(in_grad[0].asnumpy()), str(in_grad[0].shape))) |
|
|
|
|
| @mx.operator.register("identity") |
| class IdentityOpProp(mx.operator.CustomOpProp): |
| def __init__(self, logging_prefix="identity", input_debug=False, grad_debug=False): |
| super(IdentityOpProp, self).__init__(need_top_grad=True) |
| self.input_debug = safe_eval(input_debug) |
| self.grad_debug = safe_eval(grad_debug) |
| self.logging_prefix = str(logging_prefix) |
|
|
| def list_arguments(self): |
| return ['data'] |
|
|
| def list_outputs(self): |
| return ['output'] |
|
|
| def infer_shape(self, in_shape): |
| data_shape = in_shape[0] |
| output_shape = in_shape[0] |
| return [data_shape], [output_shape], [] |
|
|
| def create_operator(self, ctx, shapes, dtypes): |
| return IdentityOp(input_debug=self.input_debug, |
| grad_debug=self.grad_debug, |
| logging_prefix=self.logging_prefix) |
|
|
| class SaveNpyOp(mx.operator.CustomOp): |
| def __init__(self, save_name="op", save_dir=None): |
| super(SaveNpyOp, self).__init__() |
| self._save_name = save_name |
| self._save_dir = '.' if save_dir is None else save_dir |
| if not os.path.exists(save_dir): |
| os.makedirs(save_dir) |
| self._input_save_path = os.path.join(self._save_dir, '{}.npy'.format(save_name)) |
| self._grad_save_path = os.path.join(self._save_dir, '{}_grad.npy'.format(save_name)) |
|
|
| def forward(self, is_train, req, in_data, out_data, aux): |
| logging.info("Saving Input to {}".format(os.path.realpath(self._input_save_path))) |
| np.save(self._input_save_path, in_data[0].asnumpy()) |
| self.assign(out_data[0], req[0], in_data[0]) |
|
|
| def backward(self, req, out_grad, in_data, out_data, in_grad, aux): |
| logging.info("Saving Gradient to {}".format(os.path.realpath(self._input_save_path))) |
| np.save(self._grad_save_path, out_grad[0].asnumpy()) |
| self.assign(in_grad[0], req[0], out_grad[0]) |
|
|
|
|
| @mx.operator.register("save_npy") |
| class SaveNpyOpProp(mx.operator.CustomOpProp): |
| def __init__(self, save_name="op", save_dir="."): |
| super(SaveNpyOpProp, self).__init__(need_top_grad=True) |
| self._save_name = save_name |
| self._save_dir = save_dir |
|
|
| def list_arguments(self): |
| return ['data'] |
|
|
| def list_outputs(self): |
| return ['output'] |
|
|
| def infer_shape(self, in_shape): |
| data_shape = in_shape[0] |
| output_shape = in_shape[0] |
| return [data_shape], [output_shape], [] |
|
|
| def create_operator(self, ctx, shapes, dtypes): |
| return SaveNpyOp(save_name=self._save_name, |
| save_dir=self._save_dir) |
|
|
| class ConstantOp(mx.operator.CustomOp): |
| """Implementation of mask on minibatch layer. |
| """ |
| def __init__(self, data): |
| super(ConstantOp, self).__init__() |
| self.data = data |
|
|
| def forward(self, is_train, req, in_data, out_data, aux): |
| if self.data.context != out_data[0].context: |
| self.data = self.data.copyto(out_data[0].context) |
| self.assign(out_data[0], req[0], self.data) |
|
|
| def backward(self, req, out_grad, in_data, out_data, in_grad, aux): |
| raise RuntimeError("cannot bp to constant") |
|
|
|
|
| @mx.operator.register("constant") |
| class ConstantOpProp(mx.operator.CustomOpProp): |
| def __init__(self, pkl_data): |
| super(ConstantOpProp, self).__init__(need_top_grad=False) |
| self.data = pickle.loads(pkl_data) |
|
|
| def list_arguments(self): |
| return [] |
|
|
| def list_outputs(self): |
| return ['output'] |
|
|
| def infer_shape(self, in_shape): |
| return in_shape, [self.data.shape], [] |
|
|
| def create_operator(self, ctx, shapes, dtypes): |
| return ConstantOp(mx.nd.array(self.data, ctx=ctx)) |
|
|
|
|
| class LogisticRegressionMaskOutput(mx.operator.CustomOp): |
| def __init__(self, ignore_label): |
| super(LogisticRegressionMaskOutput, self).__init__() |
| self.ignore_label = ignore_label |
|
|
| def forward(self, is_train, req, in_data, out_data, aux): |
| self.assign(out_data[0], req[0], 1.0 / (1.0 + nd.exp(- in_data[0]))) |
|
|
| def backward(self, req, out_grad, in_data, out_data, in_grad, aux): |
| output = out_data[0].asnumpy() |
| label = in_data[1].asnumpy() |
| data_grad = (output - label) * (label != self.ignore_label) |
| self.assign(in_grad[0], req[0], data_grad) |
|
|
| @mx.operator.register("LogisticRegressionMaskOutput") |
| class LogisticRegressionMaskOutputProp(mx.operator.CustomOpProp): |
| def __init__(self, ignore_label): |
| super(LogisticRegressionMaskOutputProp, self).__init__(need_top_grad=False) |
| self.ignore_label = safe_eval(ignore_label) |
|
|
| def list_arguments(self): |
| return ['data', 'label'] |
|
|
| def list_outputs(self): |
| return ['output'] |
|
|
| def infer_shape(self, in_shape): |
| data_shape = in_shape[0] |
| label_shape = in_shape[0] |
| output_shape = in_shape[0] |
| return [data_shape, label_shape], [output_shape], [] |
|
|
| def create_operator(self, ctx, shapes, dtypes): |
| return LogisticRegressionMaskOutput(ignore_label=self.ignore_label) |
|
|
| class EntropyMultinomialDist(mx.operator.CustomOp): |
| def __init__(self): |
| super(EntropyMultinomialDist, self).__init__() |
|
|
| def forward(self, is_train, req, in_data, out_data, aux): |
| self.assign(out_data[0], req[0], scipy.stats.entropy(in_data[0].asnumpy().T)) |
|
|
| def backward(self, req, out_grad, in_data, out_data, in_grad, aux): |
| p = in_data[0] |
| p_sum = nd.sum(p, axis=1, keepdims=True) |
| logit = nd.log(p / p_sum) |
| grad = - logit / p_sum + nd.sum(p * logit, axis=1, keepdims=True) / nd.square(p_sum) |
| grad[:] = nd.expand_dims(out_grad[0], axis=1) * grad |
| self.assign(in_grad[0], req[0], grad) |
|
|
| @mx.operator.register("entropy_multinomial") |
| class EntropyMultinomialDistProp(mx.operator.CustomOpProp): |
| def __init__(self): |
| super(EntropyMultinomialDistProp, self).__init__(need_top_grad=True) |
|
|
| def list_arguments(self): |
| return ['data'] |
|
|
| def list_outputs(self): |
| return ['output'] |
|
|
| def infer_shape(self, in_shape): |
| data_shape = in_shape[0] |
| output_shape = (in_shape[0][0],) |
| return [data_shape], [output_shape], [] |
|
|
| def create_operator(self, ctx, shapes, dtypes): |
| return EntropyMultinomialDist() |
|
|
|
|
| def logistic_regression_mask_output(data, label, ignore_label, name=None): |
| return mx.sym.Custom(name=name, |
| op_type="LogisticRegressionMaskOutput", |
| ignore_label=ignore_label, |
| data=data, |
| label=label) |
|
|
|
|
| def constant(data, name="constant"): |
| if isinstance(data, mx.nd.NDArray): |
| data = data.asnumpy() |
| pkl_data = pickle.dumps(data) |
| return mx.symbol.Custom(name=name, |
| op_type="constant", |
| pkl_data=pkl_data) |
|
|
|
|
| def identity(data, name="identity", logging_prefix=None, |
| input_debug=False, grad_debug=False): |
| return mx.symbol.Custom(data=data, |
| name=name, |
| logging_prefix=name, |
| input_debug=input_debug, |
| grad_debug=grad_debug, |
| op_type="identity") |
|
|
|
|
| def save_npy(data, save_name="op", save_dir="."): |
| return mx.symbol.Custom(data=data, |
| save_name=save_name, |
| save_dir=save_dir, |
| op_type="save_npy") |
|
|
|
|
| def entropy_multinomial(data, name="entropy"): |
| return mx.symbol.Custom(name=name, |
| op_type="entropy_multinomial", |
| data=data) |
|
|
|
|
| def grid_generator(batch_size, height, width, normalize=True): |
| """Generate the grid based on width and height |
| |
| Parameters |
| ---------- |
| batch_size : int |
| width : int |
| height : int |
| normalize : bool |
| Whether to normalize the grid elements into [-1, 1] |
| |
| Returns |
| ------- |
| ret : mx.sym.Symbol |
| Shape : (batch_size, 2, height, width), the channel contains (x, y) |
| """ |
| x = mx.sym.arange(start=0, stop=width) |
| y = mx.sym.arange(start=0, stop=height) |
| x = mx.sym.broadcast_to(mx.sym.Reshape(x, shape=(1, 1, 1, width)), |
| shape=(batch_size, 1, height, width)) |
| y = mx.sym.broadcast_to(mx.sym.Reshape(y, shape=(1, 1, height, 1)), |
| shape=(batch_size, 1, height, width)) |
| if normalize: |
| x = x / float(width - 1) * 2.0 - 1.0 |
| y = y / float(height - 1) * 2.0 - 1.0 |
| ret = mx.sym.Concat(x, y, num_args=2, dim=1) |
| return ret |
|
|
|
|
| def normalize_grid(un_norm_grid, width, height): |
| """Normalize the grid to [-1, 1] |
| |
| Parameters |
| ---------- |
| un_norm_grid : mx.sym.Symbol |
| Shape : (batch_size, 2, height, width) |
| width : int |
| height : int |
| |
| Returns |
| ------- |
| ret : mx.sym.Symbol |
| """ |
| un_norm_grid = mx.sym.SliceChannel(un_norm_grid, axis=1, num_outputs=2, squeeze_axis=False) |
| x = un_norm_grid[0] / float(width - 1) * 2.0 - 1.0 |
| y = un_norm_grid[1] / float(height - 1) * 2.0 - 1.0 |
| ret = mx.sym.Concat(x, y, num_args=2, dim=1) |
| return ret |
|
|
|
|
| def multi_segment_slice_axis(data, axis, segment_lengths): |
| """Split the data to multiple segments |
| |
| Parameters |
| ---------- |
| data : mx.sym.Symbol |
| axis : int |
| segment_lengths : list or tuple |
| Get the segment_lengths |
| Returns |
| ------- |
| ret : list |
| """ |
| ret = [] |
| begin = 0 |
| for length in segment_lengths: |
| seg_ele = mx.sym.slice_axis(data=data, axis=axis, begin=begin, end=begin + length) |
| ret.append(seg_ele) |
| begin += length |
| return tuple(ret) |
|
|
|
|
| def group_add(lhs, rhs): |
| """ |
| |
| Parameters |
| ---------- |
| lhs : list of mx.sym.Symbol |
| rhs : list of mx.sym.Symbol |
| |
| Returns |
| ------- |
| ret : list of mx.sym.Symbol |
| """ |
| if isinstance(lhs, mx.sym.Symbol): |
| return lhs + rhs |
| assert len(lhs) == len(rhs) |
| ret = [] |
| for i in range(len(lhs)): |
| if isinstance(lhs[i], list): |
| ret.append(group_add(lhs[i], rhs[i])) |
| else: |
| ret.append(lhs[i] + rhs[i]) |
| return ret |
|
|
|
|
| def one_step_diff(dat, axis): |
| """ |
| |
| Parameters |
| ---------- |
| dat : mx.sym.Symbol |
| axes : tuple |
| |
| Returns |
| ------- |
| |
| """ |
| return mx.sym.slice_axis(dat, axis=axis, begin=0, end=-1) - \ |
| mx.sym.slice_axis(dat, axis=axis, begin=1, end=None) |
|
|
|
|
| def masked_gdl_loss(pred, gt, mask): |
| """ |
| |
| Parameters |
| ---------- |
| pred : mx.sym.Symbol |
| Shape: (seq_len, batch_size, 1, H, W) |
| gt : mx.sym.Symbol |
| Shape: (seq_len, batch_size, 1, H, W) |
| mask : mx.sym.Symbol |
| Shape: (seq_len, batch_size, 1, H, W) |
| |
| Returns |
| ------- |
| gdl : mx.sym.Symbol |
| Shape: (seq_len, batch_size) |
| """ |
| valid_mask_h = mx.sym.slice_axis(mask, axis=3, begin=0, end=-1) *\ |
| mx.sym.slice_axis(mask, axis=3, begin=1, end=None) |
| valid_mask_w = mx.sym.slice_axis(mask, axis=4, begin=0, end=-1) *\ |
| mx.sym.slice_axis(mask, axis=4, begin=1, end=None) |
| pred_diff_h = mx.sym.abs(one_step_diff(pred, axis=3)) |
| pred_diff_w = mx.sym.abs(one_step_diff(pred, axis=4)) |
| gt_diff_h = mx.sym.abs(one_step_diff(gt, axis=3)) |
| gt_diff_w = mx.sym.abs(one_step_diff(gt, axis=4)) |
| gd_h = mx.sym.abs(pred_diff_h - gt_diff_h) |
| gd_w = mx.sym.abs(pred_diff_w - gt_diff_w) |
| gdl = mx.sym.sum(valid_mask_h * gd_h, axis=(2, 3, 4)) +\ |
| mx.sym.sum(valid_mask_w * gd_w, axis=(2, 3, 4)) |
| return gdl |
|
|
|
|
| def weighted_l2(pred, gt, weight): |
| """ |
| |
| Parameters |
| ---------- |
| pred : mx.sym.Symbol |
| Shape: (seq_len, batch_size, 1, H, W) |
| gt : mx.sym.Symbol |
| Shape: (seq_len, batch_size, 1, H, W) |
| weight : mx.sym.Symbol |
| Shape: (seq_len, batch_size, 1, H, W) |
| |
| Returns |
| ------- |
| l2 : mx.nd.NDArray |
| Shape: (seq_len, batch_size) |
| """ |
| l2 = weight * mx.sym.square(pred - gt) |
| l2 = mx.sym.sum(l2, axis=(2, 3, 4)) |
| return l2 |
|
|
|
|
| def weighted_mse(pred, gt, weight): |
| return weighted_l2(pred, gt, weight) |
|
|
|
|
| def weighted_l1(pred, gt, weight): |
| l1 = weight * mx.sym.abs(pred - gt) |
| l1 = mx.sym.sum(l1, axis=(2, 3, 4)) |
| return l1 |
|
|
|
|
| def weighted_mae(pred, gt, weight): |
| return weighted_l1(pred, gt, weight) |
|
|
|
|
| def masked_hit_miss_counts(pred, gt, mask, thresholds): |
| """ |
| |
| Parameters |
| ---------- |
| pred : mx.sym.Symbol |
| Shape: (seq_len, batch_size, 1, H, W) |
| gt : mx.sym.Symbol |
| Shape: (seq_len, batch_size, 1, H, W) |
| mask : mx.sym.Symbol |
| Shape: (seq_len, batch_size, 1, H, W) |
| thresholds : list |
| |
| Returns |
| ------- |
| hits : mx.nd.NDArray |
| Shape: (seq_len, batch_size, len(thresholds)) |
| misses : mx.nd.NDArray |
| Shape: (seq_len, batch_size, len(thresholds)) |
| false_alarms : mx.nd.NDArray |
| Shape: (seq_len, batch_size, len(thresholds)) |
| correct_negatives : mx.nd.NDArray |
| Shape: (seq_len, batch_size, len(thresholds)) |
| """ |
| from nowcasting.hko_evaluation import rainfall_to_pixel |
| thresholds = [rainfall_to_pixel(threshold) for threshold in thresholds] |
| hits = [] |
| misses = [] |
| false_alarms = [] |
| correct_negatives = [] |
| for threshold in thresholds: |
| pred_rain_mask = pred > threshold |
| gt_rain_mask = gt > threshold |
| hits_ele = pred_rain_mask * gt_rain_mask * mask |
| misses_ele = (1 - pred_rain_mask) * gt_rain_mask * mask |
| false_alarms_ele = pred_rain_mask * (1 - gt_rain_mask) * mask |
| correct_negatives_ele = (1 - pred_rain_mask) * (1 - gt_rain_mask) * mask |
| hits.append(mx.sym.sum(hits_ele, axis=(3, 4))) |
| misses.append(mx.sym.sum(misses_ele, axis=(3, 4))) |
| false_alarms.append(mx.sym.sum(false_alarms_ele, axis=(3, 4))) |
| correct_negatives.append(mx.sym.sum(correct_negatives_ele, axis=(3, 4))) |
| hits = mx.sym.concat(*hits, dim=2, num_args=len(thresholds)) |
| misses = mx.sym.concat(*misses, dim=2, num_args=len(thresholds)) |
| false_alarms = mx.sym.concat(*false_alarms, dim=2, num_args=len(thresholds)) |
| correct_negatives = mx.sym.concat(*correct_negatives, dim=2, num_args=len(thresholds)) |
| return hits, misses, false_alarms, correct_negatives |
|
|
