def mul(g, x, y, op_scale, op_zero_point): x, _, _, _ = sym_help.dequantize_helper(g, x) y, _, _, _ = sym_help.dequantize_helper(g, y) output = mul(g, x, y) return sym_help.quantize_helper(g, output, op_scale, op_zero_point)
def group_norm_symbolic(g, input, num_groups, weight, bias, eps, cudnn_enabled): from torch.onnx.symbolic_opset9 import reshape, mul, add, reshape_as channels_num = input.type().sizes()[1] if num_groups == channels_num: output = g.op('InstanceNormalization', input, weight, bias, epsilon_f=eps) else: # Reshape from [n, g * cg, h, w] to [1, n * g, cg * h, w]. x = reshape(g, input, [0, num_groups, -1, 0]) x = reshape(g, x, [1, -1, 0, 0]) # Normalize channel-wise. x = g.op('MeanVarianceNormalization', x, axes_i=[2, 3]) # Reshape back. x = reshape_as(g, x, input) # Apply affine transform. x = mul(g, x, reshape(g, weight, [1, channels_num, 1, 1])) output = add(g, x, reshape(g, bias, [1, channels_num, 1, 1])) return output
def addcmul_symbolic(g, self, tensor1, tensor2, value=1, out=None): from torch.onnx.symbolic_opset9 import add, mul if out is not None: sym_help._unimplemented("addcmul", "Out parameter is not supported for addcmul") x = mul(g, tensor1, tensor2) value = sym_help._maybe_get_scalar(value) if sym_help._scalar(value) != 1: value = sym_help._if_scalar_type_as(g, value, x) if not sym_help._is_value(value): value = g.op("Constant", value_t=torch.tensor(value, dtype=torch.float32)) x = mul(g, x, value) return add(g, self, x)
def normal(g, loc, scale, seed): # If you can sample from a given distribution with mean 0 and variance 1, then you can easily sample from a # scale-location transformation of that distribution, which has mean μ and variance σ's square. If x is a sample # from a mean 0 and variance 1 distribution then # σx+μ # is a sample with mean μ and variance σ's square. result = mul(g, scale, g.op("RandomNormalLike", loc)) return add(g, result, loc)
def mul(g, x, y, op_scale, op_zero_point): x, _, _, _ = symbolic_helper.dequantize_helper(g, x) y, _, _, _ = symbolic_helper.dequantize_helper(g, y) output = opset9.mul(g, x, y) return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point)
def binary_cross_entropy_with_logits(g, input, target, weight, pos_weight, reduction): from torch.onnx.symbolic_opset9 import sigmoid, log, sub, neg, mul, add p = g.op("Constant", value_t=torch.tensor([1])) sig_x = sigmoid(g, input) log_sig_x = log(g, sig_x) sub_1_x = sub(g, p, sig_x) sub_1_y = sub(g, p, target) log_1_x = log(g, sub_1_x) if pos_weight is None or sym_help._is_none(pos_weight): output = neg( g, add(g, mul(g, target, log_sig_x), mul(g, sub_1_y, log_1_x))) else: output = neg( g, add(g, mul(g, mul(g, target, log_sig_x), pos_weight), mul(g, sub_1_y, log_1_x))) if weight is not None and not sym_help._is_none(weight): output = mul(g, weight, output) reduction = sym_help._maybe_get_const(reduction, 'i') if reduction == 0: return output elif reduction == 1: return g.op("ReduceMean", output) elif reduction == 2: return g.op("ReduceSum", output) else: return sym_help._onnx_unsupported( "binary_cross_entropy_with_logits with reduction other than none, mean, or sum" )
def binary_cross_entropy_with_logits(g, input, target, weight, pos_weight, reduction): p = g.op("Constant", value_t=torch.tensor([1])) sig_x = opset9.sigmoid(g, input) log_sig_x = opset9.log(g, sig_x) sub_1_x = opset9.sub(g, p, sig_x) sub_1_y = opset9.sub(g, p, target) log_1_x = opset9.log(g, sub_1_x) if pos_weight is None or symbolic_helper._is_none(pos_weight): output = opset9.neg( g, opset9.add(g, opset9.mul(g, target, log_sig_x), opset9.mul(g, sub_1_y, log_1_x)), ) else: output = opset9.neg( g, opset9.add( g, opset9.mul(g, opset9.mul(g, target, log_sig_x), pos_weight), opset9.mul(g, sub_1_y, log_1_x), ), ) if weight is not None and not symbolic_helper._is_none(weight): output = opset9.mul(g, weight, output) reduction = symbolic_helper._maybe_get_const(reduction, "i") if reduction == 0: return output elif reduction == 1: return g.op("ReduceMean", output, keepdims_i=0) elif reduction == 2: return g.op("ReduceSum", output, keepdims_i=0) else: return symbolic_helper._onnx_unsupported( "binary_cross_entropy_with_logits with reduction other than none, mean, or sum", input, )
def multiclass_nms_core_symbolic(g, multi_bboxes, multi_scores, score_thr, nms_cfg, max_num=-1): from torch.onnx.symbolic_opset9 import reshape, squeeze from torch.onnx.symbolic_opset10 import _slice def cast(x, dtype): return g.op('Cast', x, to_i=sym_help.cast_pytorch_to_onnx[dtype]) def get_size(x, dim): shape = g.op('Shape', x) dim = _slice(g, shape, axes=[0], starts=[dim], ends=[dim + 1]) return cast(dim, 'Long') nms_op_type = nms_cfg.get('type', 'nms') assert nms_op_type == 'nms' assert 'iou_thr' in nms_cfg iou_threshold = nms_cfg['iou_thr'] assert 0 <= iou_threshold <= 1 # Transpose and reshape input tensors to fit ONNX NonMaxSuppression. multi_bboxes = reshape(g, multi_bboxes, [0, -1, 4]) multi_bboxes = g.op('Transpose', multi_bboxes, perm_i=[1, 0, 2]) batches_num = get_size(multi_bboxes, 0) spatial_num = get_size(multi_bboxes, 1) multi_scores = g.op('Transpose', multi_scores, perm_i=[1, 0]) scores_shape = g.op('Concat', batches_num, g.op('Constant', value_t=torch.LongTensor([-1])), spatial_num, axis_i=0) multi_scores = reshape(g, multi_scores, scores_shape) classes_num = get_size(multi_scores, 1) assert max_num > 0 indices = g.op( 'NonMaxSuppression', multi_bboxes, multi_scores, g.op('Constant', value_t=torch.LongTensor([max_num])), g.op('Constant', value_t=torch.FloatTensor([iou_threshold])), g.op('Constant', value_t=torch.FloatTensor([score_thr]))) # Flatten bboxes and scores. multi_bboxes_flat = reshape(g, multi_bboxes, [-1, 4]) multi_scores_flat = reshape(g, multi_scores, [ -1, ]) # Flatten indices. batch_indices = _slice(g, indices, axes=[1], starts=[0], ends=[1]) class_indices = _slice(g, indices, axes=[1], starts=[1], ends=[2]) box_indices = _slice(g, indices, axes=[1], starts=[2], ends=[3]) def add(*args, dtype='Long'): x = g.op('Add', args[0], args[1]) if dtype is not None: x = cast(x, dtype) return x def mul(*args, dtype='Long'): x = g.op('Mul', args[0], args[1]) if dtype is not None: x = cast(x, dtype) return x flat_box_indices = add(mul(batch_indices, spatial_num), box_indices) flat_score_indices = add( mul(add(mul(batch_indices, classes_num), class_indices), spatial_num), box_indices) # Select bboxes. out_bboxes = reshape( g, g.op('Gather', multi_bboxes_flat, flat_box_indices, axis_i=0), [-1, 4]) out_scores = reshape( g, g.op('Gather', multi_scores_flat, flat_score_indices, axis_i=0), [-1, 1]) # Having either batch size or number of classes here equal to one is the limitation of implementation. class_indices = reshape(g, cast(add(class_indices, batch_indices), 'Float'), [-1, 1]) # Combine bboxes, scores and labels into a single tensor. # This a workaround for a PyTorch bug (feature?), # limiting ONNX operations to output only single tensor. out_combined_bboxes = g.op('Concat', out_bboxes, out_scores, class_indices, axis_i=1) # Get the top scored bboxes only. elements_num = sym_help._size_helper(g, out_scores, dim=g.op('Constant', value_t=torch.LongTensor( [0]))) max_num = g.op('Constant', value_t=torch.LongTensor([max_num])) if sym_help._export_onnx_opset_version < 12: kn = g.op('Concat', max_num, elements_num, axis_i=0) kn = g.op('ReduceMin', kn, keepdims_i=0) else: kn = g.op('Min', max_num, elements_num) _, top_indices = sym_help._topk_helper(g, out_scores, kn, dim=0) # top_indices = squeeze(g, top_indices, dim=1) top_indices = reshape(g, top_indices, [ -1, ]) out_combined_bboxes = g.op('Gather', out_combined_bboxes, top_indices, axis_i=0) return out_combined_bboxes