def top_k_exporter_v1(op_def, context): node, (k, axis, largest, sorted) = top_k_exporter(**locals()) if largest == 0: raise ValueError('TopK-1 does not support smallest mode.') helper.add_attribute(node, 'axis', axis) helper.add_attribute(node, 'k', k) return node, None
def multinomial_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) helper.add_attribute(node, 'dtype', helper.tensor_type('int64')) for arg in op_def.arg: if arg.name == 'sample_size': helper.add_attribute(node, 'sample_size', arg.i) return node, const_tensors
def depth_space_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) for arg in op_def.arg: _assert_data_format(arg) if arg.name == 'block_size': helper.add_attribute(node, 'blocksize', arg.i) return node, const_tensors
def pool(op_def, context): node, const_tensors = export_util.translate(**locals()) rank = len(context.blob_shapes[op_def.input[0]]) - 2 global_pool, node_copy = 0, copy.deepcopy(node) for arg in op_def.arg: _assert_data_format(arg) if arg.name == 'kernel_shape': helper.add_attribute(node, 'kernel_shape', _normalize_tuple(arg.ints, rank)) elif arg.name == 'strides': helper.add_attribute(node, 'strides', _normalize_tuple(arg.ints, rank)) elif arg.name == 'pads': helper.add_attribute(node, 'pads', _normalize_pads(arg.ints, rank)) elif arg.name == 'padding' and arg.s != b'VALID': helper.add_attribute(node, 'auto_pad', arg.s) elif arg.name == 'mode': if arg.s == b'MAX': node.op_type = 'MaxPool' elif arg.s == b'AVG': node.op_type = 'AveragePool' elif arg.name == 'ceil_mode': helper.add_attribute(node, 'ceil_mode', arg.i) elif arg.name == 'global_pool': global_pool = arg.i if global_pool > 0: # Remove regular pooling attributes. node_copy.op_type = 'Global' + node.op_type node = node_copy return node, const_tensors
def affine_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) node.op_type = 'ATen' # Currently not supported in ai.onnx helper.add_attribute(node, 'op_type', 'Affine') for arg in op_def.arg: if arg.name == 'axes': helper.add_attribute(node, 'axes', arg.ints) return node, const_tensors
def clip_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) for arg in op_def.arg: if arg.name == 'low': helper.add_attribute(node, 'min', arg.f) elif arg.name == 'high': helper.add_attribute(node, 'max', arg.f) return node, const_tensors
def relu_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) for arg in op_def.arg: if arg.name == 'alpha': if arg.f > 0: node.op_type = 'LeakyRelu' helper.add_attribute(node, 'alpha', arg.f) return node, const_tensors
def softmax_exporter_v13(op_def, context): node, const_tensors = export_util.translate(**locals()) ndim = len(context.blob_shapes[op_def.input[0]]) for arg in op_def.arg: if arg.name == 'axis': axis = arg.i + (ndim if arg.i < 0 else 0) helper.add_attribute(node, 'axis', arg.i) return node, const_tensors
def unsqueeze_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) axes = None for arg in op_def.arg: if arg.name == 'axes': axes = arg.ints if axes is not None: helper.add_attribute(node, 'axes', axes) return node, const_tensors
def transpose_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) for arg in op_def.arg: if arg.name == 'perm': helper.add_attribute(node, 'perm', arg.ints) elif arg.name == 'perm_desc': values = helper.fetch_argument(op_def, arg, context.ws) helper.add_attribute(node, 'perm', values) return node, None
def cast_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) node.op_type = 'Cast' if len(node.input) == 0: raise ValueError('ONNX does not support in-place cast.') for arg in op_def.arg: if arg.name == 'dtype': helper.add_attribute(node, 'to', helper.tensor_type(arg.s)) return node, const_tensors
def split_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) axis = 0 for arg in op_def.arg: if arg.name == 'axis': axis = arg.i size_splits = [context.blob_shapes[e][axis] for e in op_def.output] helper.add_attribute(node, 'split', size_splits) return node, const_tensors
def lp_normalize_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) node.op_type = 'LpNormalization' for arg in op_def.arg: if arg.name == 'axis': helper.add_attribute(node, 'axis', arg.i) if arg.name == 'p': helper.add_attribute(node, 'p', arg.i) return node, const_tensors
def reduce_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) axes = list(range(len(context.blob_shapes[op_def.input[0]]))) for arg in op_def.arg: if arg.name == 'axes': axes = arg.ints elif arg.name == 'keepdims': helper.add_attribute(node, 'keepdims', arg.i) helper.add_attribute(node, 'axes', axes) return node, const_tensors
def fill_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) node.op_type = 'Constant' shape = list(context.blob_shapes[op_def.output[0]]) value = helper.from_array( numpy.array(shape, 'int64'), context.unique_name(op_def.output[0] + '/constant/value')) helper.add_attribute(node, 'value', value) node.ClearField('input') return node, [value]
def index_select_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) node.op_type = 'Gather' for arg in op_def.arg: if arg.name == 'axis': helper.add_attribute(node, 'axis', arg.i) elif arg.name == 'num_axes': if arg.i > 1: raise ValueError('Reshape to avoid selecting multiple axes.') return node, const_tensors
def dropout_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) drop_ratio = 0.5 # The prob to set zeros randomly. for arg in op_def.arg: if arg.name == 'prob': drop_ratio = arg.f elif arg.name == 'prob_desc': drop_ratio = helper.fetch_argument(op_def, arg, context.ws) helper.add_attribute(node, 'ratio', float(drop_ratio)) return node, const_tensors
def softmax_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) ndim = len(context.blob_shapes[op_def.input[0]]) for arg in op_def.arg: if arg.name == 'axis': axis = arg.i + (ndim if arg.i < 0 else 0) if axis != (ndim - 1): raise ValueError('Axis could only be the last if opset < 13.') helper.add_attribute(node, 'axis', arg.i) return node, const_tensors
def selu_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) alpha, gamma = 1.67326, 1.0507 for arg in op_def.arg: if arg.name == 'alpha': alpha = arg.f elif arg.name == 'gamma': gamma = arg.f helper.add_attribute(node, 'alpha', alpha) helper.add_attribute(node, 'gamma', gamma) return node, const_tensors
def flatten_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) for arg in op_def.arg: if arg.name == 'axis': helper.add_attribute(node, 'axis', arg.i) elif arg.name == 'num_axes': if arg.i != -1: raise ValueError( 'Excepted <num_axes> is -1, ' 'got {}.'.format(arg.i)) return node, None
def hardsigmoid_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) alpha, beta = 0.2, 0.5 for arg in op_def.arg: if arg.name == 'alpha': alpha = arg.f elif arg.name == 'beta': beta = arg.f helper.add_attribute(node, 'alpha', alpha) helper.add_attribute(node, 'beta', beta) return node, const_tensors
def top_k_exporter_v10(op_def, context): node, (k, axis, largest, sorted) = top_k_exporter(**locals()) if largest == 0: raise ValueError('TopK-10 does not support smallest mode.') helper.add_attribute(node, 'axis', axis) k = helper.from_array( numpy.array([k], 'int64'), context.unique_name(op_def.input[0] + '/top_k/k'), ) node.input.extend([k.name]) return node, [k]
def arg_reduce_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) # ONNX requires indices only, remove the values. indices = node.output[0] node.ClearField('output') node.output.extend([indices]) for arg in op_def.arg: if arg.name == 'axis': helper.add_attribute(node, 'axis', arg.i) elif arg.name == 'keepdims': helper.add_attribute(node, 'keepdims', arg.i) return node, None
def roi_pool(op_def, context): node, const_tensors = export_util.translate(**locals()) node.op_type = 'MaxRoiPool' pooled_shape = [None, None] for arg in op_def.arg: if arg.name == 'pooled_h': pooled_shape[0] = arg.i elif arg.name == 'pooled_w': pooled_shape[1] = arg.i elif arg.name == 'spatial_scale': helper.add_attribute(node, 'spatial_scale', arg.f) helper.add_attribute(node, 'pooled_shape', pooled_shape) return node, const_tensors
def resize_v7(op_def, context): node, const_tensors = export_util.translate(**locals()) node.op_type = 'Upsample' input_shape = context.blob_shapes[op_def.input[0]] output_shape = context.blob_shapes[op_def.output[0]] for arg in op_def.arg: if arg.name == 'mode': helper.add_attribute(node, 'mode', arg.s.lower()) helper.add_attribute(node, 'scales', [ float(output_shape[i]) / input_shape[i] for i in range(len(input_shape)) ]) return node, const_tensors
def pad_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) pads, value = [], 0 for arg in op_def.arg: if arg.name == 'pads': pads = [int(e) for e in arg.ints] elif arg.name == 'pads_desc': pads = helper.fetch_argument(op_def, arg, context.ws) elif arg.name == 'mode': helper.add_attribute(node, 'mode', arg.s.lower()) elif arg.name == 'value': value = arg.f return node, pads, value
def trilu_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) k = 0 for arg in op_def.arg: if arg.name == 'upper': helper.add_attribute(node, 'upper', arg.i) elif arg.name == 'k': k = arg.i k = helper.from_array( numpy.array(k, 'int64'), context.unique_name(op_def.input[0] + '/trilu/k'), ) node.input.extend([k.name]) return node, [k]
def gather_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) axis, end_axis = None, None for arg in op_def.arg: if arg.name == 'axis': axis = arg.i helper.add_attribute(node, 'axis', arg.i) elif arg.name == 'end_axis': end_axis = arg.i if end_axis < 0: input_shape = context.blob_shapes[op_def.input[0]] end_axis += len(input_shape) if end_axis is not None and axis != end_axis: raise ValueError('Reshape to avoid multiple axes.') return node, const_tensors
def cumulative_exporter(op_def, context): node, const_tensors = export_util.translate(**locals()) axis = 0 for arg in op_def.arg: if arg.name == 'axis': axis = arg.i elif arg.name == 'exclusive': helper.add_attribute(node, 'exclusive', arg.i) elif arg.name == 'reverse': helper.add_attribute(node, 'reverse', arg.i) axis = helper.from_array( numpy.array(axis, 'int64'), context.unique_name(op_def.input[0] + '/cumulative/axis'), ) node.input.extend([axis.name]) return node, [axis]
def resize_v10(op_def, context): node, const_tensors = export_util.translate(**locals()) input_shape = context.blob_shapes[op_def.input[0]] output_shape = context.blob_shapes[op_def.output[0]] for arg in op_def.arg: if arg.name == 'mode': helper.add_attribute(node, 'mode', arg.s.lower()) scales = helper.from_array( numpy.array([ float(output_shape[i]) / input_shape[i] for i in range(len(input_shape)) ], 'float32'), context.unique_name(op_def.input[0] + '/resize/scales'), ) node.input.extend([scales.name]) return node, [scales]