def unfold(g, input, dimension, size, step):
const_size = sym_help._maybe_get_const(size, "i")
const_step = sym_help._maybe_get_const(step, "i")
if not sym_help._is_value(const_size) and not sym_help._is_value(const_step):
from torch.onnx.symbolic_opset9 import unfold as _unfold
return _unfold(g, input, dimension, const_size, const_step)
if sym_help._operator_export_type == torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK:
return g.op("ATen", input, operator_s="unfold", dimension_i=dimension, size_i=size, step_i=step)
sizedim = sym_help._get_tensor_dim_size(input, dimension)
if sizedim is not None:
low_start = g.op("Constant", value_t=torch.tensor(0))
low_end = g.op("Constant", value_t=torch.tensor(sizedim))
hi_end = g.op("Constant", value_t=torch.tensor(sizedim + 1))
low_indices = g.op("Range", low_start, low_end, step)
hi_indices = g.op("Range", size, hi_end, step)
low_size = sym_help._size_helper(g, low_indices, g.op("Constant", value_t=torch.tensor(0)))
hi_size = sym_help._size_helper(g, hi_indices, g.op("Constant", value_t=torch.tensor(0)))
ndim = sym_help._get_tensor_rank(input)
perm = list(range(0, ndim))
perm.append(perm.pop(dimension))
unsqueeze_list = []
loop_condition = g.op("Constant", value_t=torch.tensor(1))
loop_condition = g.op("Cast", loop_condition, to_i=9)
loop_len = g.op("Min", low_size, hi_size)
loop = g.op("Loop", loop_len, loop_condition)
loop_block = _add_block(loop.node())
block_input_iter = _add_input_to_block(loop_block)
cond = _add_input_to_block(loop_block)
starts = loop_block.op("Gather", low_indices, block_input_iter)
ends = loop_block.op("Gather", hi_indices, block_input_iter)
axes = loop_block.op("Constant", value_t=torch.tensor([2]))
starts = sym_help._unsqueeze_helper(loop_block, starts, [0])
ends = sym_help._unsqueeze_helper(loop_block, ends, [0])
stack = loop_block.op("Slice", input, starts, ends, axes)
unsqueeze = sym_help._unsqueeze_helper(loop_block, loop_block.op("Transpose", stack, perm_i=perm), [dimension])
unsqueeze_list.append(unsqueeze)
concat = loop_block.op("Concat", *unsqueeze_list, axis_i=0)
cond_out = loop_block.op("Cast", loop_condition, to_i=9)
_add_output_to_block(loop_block, cond_out)
_add_output_to_block(loop_block, concat)
loop_output = loop.node().output()
perm = [0, 1, 2, 3, 4]
perm[0], perm[dimension + 1] = perm[dimension + 1], perm[0]
transpose = g.op("Transpose", loop_output, perm_i=perm)
squeeze = sym_help._squeeze_helper(g, transpose, [0])
return squeeze
else:
return _unimplemented("Unfold", "input size not accessible")
def __interpolate(g, input, size, scale_factor, mode, align_corners, recompute_scale_factor):
align_corners = sym_help._maybe_get_const(align_corners, 'b')
if not sym_help._is_none(align_corners) and align_corners:
return _unimplemented("interpolate", "align_corners == True")
if not sym_help._is_none(scale_factor) and sym_help._is_value(scale_factor):
return _unimplemented("interpolate", "dynamic scales in opset 8")
if not sym_help._is_none(size) and sym_help._is_value(size):
return _unimplemented("interpolate", "dynamic size in opset 8")
scales, mode = sym_help._interpolate_get_scales_and_mode(g, input, size, scale_factor,
mode , align_corners)
return g.op("Upsample", input, mode_s=mode, scales_f=scales)
def scatter_add(g, self, dim, index, src):
if symbolic_helper.is_caffe2_aten_fallback():
return g.at("scatter", self, dim, index, src, overload_name="src")
src_type = src.type().scalarType()
src_sizes = symbolic_helper._get_tensor_sizes(src)
index_sizes = symbolic_helper._get_tensor_sizes(index)
if src_sizes != index_sizes:
return symbolic_helper._unimplemented(
"scatter_add",
f"`index` ({index_sizes}) should have the same dimensionality as `src` ({src_sizes})",
)
src = symbolic_helper._maybe_get_scalar(src)
if symbolic_helper._is_value(src):
return g.op("ScatterElements", self, index, src, axis_i=dim, reduction_s="add")
else:
# Check if scalar "src" has same type as self (PyTorch allows different
# type for scalar src (but not when src is tensor)). If not, insert Cast node.
if self.type().scalarType() != src_type:
src = g.op(
"Cast",
src,
to_i=symbolic_helper.cast_pytorch_to_onnx[self.type().scalarType()],
)
return g.op(
"ScatterElements",
self,
index,
src,
axis_i=dim,
reduction_s="add",
)
def frobenius_norm(g, self, dim=None, keepdim=False):
dim_val = sym_help._maybe_get_const(dim, "is")
if not sym_help._is_value(dim_val) and len(dim_val) == 0:
return g.op("ReduceL2", self, keepdims_i=0)
sqr = g.op("Mul", self, self)
sumsqr = sym_help._reducesum_helper(g, sqr, dim, keepdims_i=keepdim)
return g.op("Sqrt", sumsqr)
def frobenius_norm(g, self, dim=None, keepdim=False):
dim_val = sym_help._maybe_get_const(dim, 'is')
if not sym_help._is_value(dim_val) and len(dim_val) == 0:
return g.op("ReduceL2", self, keepdims_i=0)
sqr = g.op('Mul', self, self)
sumsqr = g.op('ReduceSum', sqr, dim, keepdims_i=keepdim)
return g.op('Sqrt', sumsqr)
def split_with_sizes(g, self, split_sizes, dim):
if sym_help._is_value(split_sizes) and split_sizes.node().kind(
) == 'prim::ListConstruct':
return g.op("SplitToSequence", self, split_sizes, axis_i=dim)
else:
return torch.onnx.symbolic_opset9.split_with_sizes(
g, self, split_sizes, dim)
def symbolic_fn(g, input, output_size, align_corners=None):
if align_corners:
return _unimplemented(name, "align_corners == True")
output_size = sym_help._maybe_get_const(output_size, 'is')
if sym_help._is_value(output_size):
offset = 2
offsets = g.op("Constant",
value_t=torch.tensor([1. for i in range(offset)]))
dividend = g.op("Cast",
output_size,
to_i=sym_help.cast_pytorch_to_onnx["Float"])
divisor = sym_help._slice_helper(g,
g.op("Shape", input),
axes=[0],
ends=[dim],
starts=[offset])
divisor = g.op("Cast",
divisor,
to_i=sym_help.cast_pytorch_to_onnx["Float"])
scale_dims = g.op("Div", dividend, divisor)
scales = g.op("Concat", offsets, scale_dims, axis_i=0)
else:
scales_constant = [
1. if i < 2 else float(output_size[-(dim - i)]) /
float(input.type().sizes()[-(dim - i)]) for i in range(0, dim)
]
scales = g.op("Constant", value_t=torch.tensor(scales_constant))
return g.op("Resize", input, scales, mode_s=interpolate_mode)
def topk_symbolic(g, self, k, dim, largest, sorted, out=None):
def reverse(x):
from torch.onnx.symbolic_opset9 import reshape, transpose, size
y = transpose(g, x, 0, dim)
shape = g.op("Shape", y)
y = reshape(g, y, [0, 1, -1])
n = size(g, y, g.op("Constant", value_t=torch.LongTensor([0])))
y = g.op("ReverseSequence", y, n, batch_axis_i=1, time_axis_i=0)
y = reshape(g, y, shape)
y = transpose(g, y, 0, dim)
return y
if out is not None:
sym_help._unimplemented("TopK",
"Out parameter is not supported for topk")
k = sym_help._maybe_get_const(k, 'i')
if not sym_help._is_value(k):
k = g.op("Constant", value_t=torch.tensor(k, dtype=torch.int64))
from torch.onnx.symbolic_opset9 import unsqueeze
k = unsqueeze(g, k, 0)
top_values, top_indices = g.op("TopK", self, k, axis_i=dim, outputs=2)
if not largest:
top_values = reverse(top_values)
top_indices = reverse(top_indices)
return top_values, top_indices
def upsample_nearest2d(g, input, output_size):
output_size = _maybe_get_const(output_size, 'is')
if _is_value(output_size):
div_lhs = g.op('Cast', output_size, to_i=cast_pytorch_to_onnx['Float'])
div_rhs = g.op('Cast',
g.op(
'Slice', g.op('Shape', input),
g.op('Constant',
value_t=torch.tensor([2], dtype=torch.long)),
g.op('Constant',
value_t=torch.tensor([4], dtype=torch.long))),
to_i=cast_pytorch_to_onnx['Float'])
scales = g.op('Concat',
g.op('Constant', value_t=torch.tensor([1., 1.])),
g.op('Div', div_lhs, div_rhs),
axis_i=0)
else:
height_scale = float(output_size[-2]) / input.type().sizes()[-2]
width_scale = float(output_size[-1]) / input.type().sizes()[-1]
scales = g.op("Constant",
value_t=torch.tensor([1., 1., height_scale,
width_scale]))
return g.op(
"Resize",
input,
scales, #'Upsample' for opset 9
mode_s="nearest")
def topk_symbolic(g, self, k, dim, largest, sorted, out=None):
from torch.onnx.symbolic_opset9 import unsqueeze
def reverse(x):
from torch.onnx.symbolic_opset9 import reshape, transpose, size
y = transpose(g, x, 0, dim)
shape = g.op("Shape", y)
y = reshape(g, y, [0, 1, -1])
n = size(g, y, g.op("Constant", value_t=torch.LongTensor([0])))
y = g.op("ReverseSequence", y, n, batch_axis_i=1, time_axis_i=0)
y = reshape(g, y, shape)
y = transpose(g, y, 0, dim)
return y
k = sym_help._maybe_get_const(k, 'i')
if not sym_help._is_value(k):
k = g.op("Constant", value_t=torch.tensor(k, dtype=torch.int64))
k = unsqueeze(g, k, 0)
do_reverse = False
if sym_help._export_onnx_opset_version <= 10 and not largest:
do_reverse = True
largest = True
top_values, top_indices = sym_help._topk_helper(g, self, k, dim, largest,
sorted, out)
if sym_help._export_onnx_opset_version <= 10 and do_reverse:
top_values = reverse(top_values)
top_indices = reverse(top_indices)
return top_values, top_indices
def symbolic_fn(g, input, output_size, align_corners=None):
align_corners = sym_help._maybe_get_scalar(align_corners)
output_size = sym_help._maybe_get_const(output_size, 'is')
if sym_help._is_value(output_size):
offsets = g.op("Constant",
value_t=torch.ones(2, dtype=torch.int64))
output_size = g.op("Cast",
output_size,
to_i=sym_help.cast_pytorch_to_onnx["Long"])
output_size = g.op("Concat", offsets, output_size, axis_i=0)
else:
output_size = [
1 if i < 2 else output_size[-(dim - i)] for i in range(0, dim)
]
output_size = g.op("Constant", value_t=torch.tensor(output_size))
coordinate_transformation_mode = "asymmetric" if interpolate_mode == "nearest" \
else "align_corners" if align_corners else "pytorch_half_pixel"
empty_tensor = g.op("Constant",
value_t=torch.tensor([], dtype=torch.float32))
return g.op(
"Resize",
input,
empty_tensor, # roi only takes effect whith coordinate_transformation_mode="tf_crop_and_resize"
empty_tensor, # scales is not needed since we are sending out_size
output_size,
coordinate_transformation_mode_s=coordinate_transformation_mode,
cubic_coeff_a_f=-0.75, # only valid when mode="cubic"
mode_s=interpolate_mode, # nearest, linear, or cubic
nearest_mode_s="floor") # only valid when mode="nearest"
def full(g, sizes, value, dtype, layout, device, pin_memory=False):
const_value = sym_help._maybe_get_const(value, 't')
if sym_help._is_value(const_value):
tmp = zeros(g, sizes, dtype, layout, device)
return sym_opset9.add(g, tmp, value, g.op("Constant", value_t=torch.tensor(1)))
else:
dtype = sym_help._get_const(dtype, 'i', 'dtype')
return _constant_fill(g, sizes, dtype, const_value)
def full(g, sizes, value, dtype, layout, device, pin_memory=False):
const_value = symbolic_helper._maybe_get_const(value, "t")
if symbolic_helper._is_value(const_value):
tmp = zeros(g, sizes, dtype, layout, device)
return opset9.add(g, tmp, value, g.op("Constant", value_t=torch.tensor(1)))
else:
dtype = symbolic_helper._get_const(dtype, "i", "dtype")
return _constant_fill(g, sizes, dtype, const_value)
def topk(g, self, k, dim, largest, sorted, out=None):
if out is not None:
_unimplemented("TopK", "Out parameter is not supported for topk")
if not largest:
_unimplemented("TopK", "Ascending TopK is not supported")
k = sym_help._maybe_get_const(k, 'i')
if not sym_help._is_value(k):
k = g.op("Constant", value_t=torch.tensor(k, dtype=torch.int64))
from torch.onnx.symbolic_opset9 import unsqueeze
k = unsqueeze(g, k, 0)
return g.op("TopK", self, k, axis_i=dim, outputs=2)
def view(g, self, size):
size = sym_help._maybe_get_const(size, 'is')
if sym_help._is_value(size):
shape = size
else:
if self.isCompleteTensor():
self_sizes = self.type().sizes()
if self_sizes and len(size) == 2 and self_sizes[0] == size[0]:
old_type, self = _try_cast_integer_to_float(g, self)
return _cast_to_type(g, g.op("Flatten", self, axis_i=1), old_type)
shape = g.op("Constant", value_t=torch.LongTensor(size))
return g.op("Reshape", self, shape)
def add(g, self, other, alpha=None):
if sym_help._is_value(self) and sym_help._is_tensor_list(self):
tensor_list_node = other.node()
if tensor_list_node.kind() != "prim::ListConstruct":
return _unimplemented("add", "does not support adding dynamic tensor list to another")
tensors = sym_help._unpack_list(other)
l = self
for t in tensors:
l = g.op("SequenceInsert", l, t)
return l
return torch.onnx.symbolic_opset9.add(g, self, other, alpha)
def symbolic_fn(g, input, output_size, align_corners=None):
sym_help._interpolate_warning(interpolate_mode)
align_corners = sym_help._maybe_get_scalar(align_corners)
if align_corners:
return _unimplemented(name, "align_corners == True")
output_size = sym_help._maybe_get_const(output_size, 'is')
if sym_help._is_value(output_size):
return _unimplemented(name, "torch._C.Value (output_size) indexing")
else:
scales = [1. if i < 2 else
float(output_size[-(dim - i)]) / float(input.type().sizes()[-(dim - i)])
for i in range(0, dim)]
return g.op("Upsample", input, mode_s=interpolate_mode, scales_f=scales)
def scatter(g, self, dim, index, src):
from torch.onnx.symbolic_opset9 import expand_as
if sym_help._operator_export_type == torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK:
return g.op("ATen", self, dim, index, src, operator_s="scatter")
src = sym_help._maybe_get_scalar(src)
if sym_help._is_value(src):
return g.op("ScatterElements", self, index, src, axis_i=dim)
else:
return g.op("ScatterElements",
self,
index,
expand_as(g, src, index),
axis_i=dim)
def repeat(g, self, repeats):
if not sym_help._is_value(repeats):
repeats = g.op("Constant", value_t=torch.LongTensor(repeats))
if sym_help._is_packed_list(repeats):
repeat_size_len = len(sym_help._unpack_list(repeats))
else:
const_repeats = sym_help._maybe_get_const(repeats, 'is')
repeat_size_len = len(const_repeats)
if self.isCompleteTensor():
sizes = self.type().sizes()
diff_dims = repeat_size_len - len(sizes)
if diff_dims > 0:
self = sym_opset9.view(g, self, [1] * diff_dims + sizes)
return g.op("Tile", self, repeats)
def scatter(g, self, dim, index, src):
from torch.onnx.symbolic_opset9 import expand_as
if sym_help._operator_export_type == torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK:
return g.op("ATen", self, dim, index, src, operator_s="scatter")
src_type = src.type().scalarType()
src = sym_help._maybe_get_scalar(src)
if sym_help._is_value(src):
return g.op("ScatterElements", self, index, src, axis_i=dim)
else:
# Check if scalar "src" has same type as self (PyTorch allows different
# type for scalar src (but not when src is tensor)). If not, insert Cast node.
if self.type().scalarType() != src_type:
src = g.op("Cast", src, to_i=sym_help.cast_pytorch_to_onnx[self.type().scalarType()])
return g.op("ScatterElements", self, index, expand_as(g, src, index), axis_i=dim)
def upsample_nearest2d(g, input, output_size, align_corners=None):
align_corners = sym_help._maybe_get_scalar(align_corners)
if align_corners:
return _unimplemented("upsample_neareset2d", "align_corners == True")
output_size = sym_help._maybe_get_const(output_size, 'is')
if sym_help._is_value(output_size):
return _unimplemented("upsample_nearest2d",
"torch._C.Value (output_size) indexing")
else:
height_scale = float(output_size[-2]) / input.type().sizes()[-2]
width_scale = float(output_size[-1]) / input.type().sizes()[-1]
scales = [1., 1., height_scale, width_scale]
return g.op("Upsample", input, mode_s="nearest", scales_f=scales)
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 scatter(g, self, dim, index, src):
if symbolic_helper.is_caffe2_aten_fallback():
return g.at("scatter", self, dim, index, src, overload_name="src")
src_type = src.type().scalarType()
src = symbolic_helper._maybe_get_scalar(src)
if symbolic_helper._is_value(src):
return g.op("ScatterElements", self, index, src, axis_i=dim)
else:
# Check if scalar "src" has same type as self (PyTorch allows different
# type for scalar src (but not when src is tensor)). If not, insert Cast node.
if self.type().scalarType() != src_type:
src = g.op(
"Cast",
src,
to_i=symbolic_helper.cast_pytorch_to_onnx[self.type().scalarType()],
)
return g.op(
"ScatterElements", self, index, opset9.expand_as(g, src, index), axis_i=dim
)
def symbolic_fn(g, input, output_size, *args):
scales, align_corners = symbolic_helper._get_interpolate_attributes(
g, interpolate_mode, args
)
symbolic_helper._interpolate_warning(interpolate_mode)
align_corners = symbolic_helper._maybe_get_scalar(align_corners)
if align_corners:
return symbolic_helper._unimplemented(name, "align_corners == True", input)
output_size = symbolic_helper._maybe_get_const(output_size, "is")
if symbolic_helper._is_value(output_size):
return symbolic_helper._unimplemented(
name, "torch._C.Value (output_size) indexing"
)
if scales is None:
scales = [
1.0
if i < 2
else float(output_size[-(dim - i)])
/ float(input.type().sizes()[-(dim - i)])
for i in range(0, dim)
]
return g.op("Upsample", input, mode_s=interpolate_mode, scales_f=scales)
def unfold(g, input, dimension, size, step):
const_size = symbolic_helper._maybe_get_const(size, "i")
const_step = symbolic_helper._maybe_get_const(step, "i")
if not symbolic_helper._is_value(
const_size) and not symbolic_helper._is_value(const_step):
return opset9.unfold(g, input, dimension, const_size, const_step)
if symbolic_helper.is_caffe2_aten_fallback():
return g.at("unfold",
input,
dimension_i=dimension,
size_i=size,
step_i=step)
sizedim = symbolic_helper._get_tensor_dim_size(input, dimension)
if sizedim is not None:
low_start = g.op("Constant", value_t=torch.tensor(0))
low_end = g.op("Constant", value_t=torch.tensor(sizedim))
hi_end = g.op("Constant", value_t=torch.tensor(sizedim + 1))
low_indices = g.op("Range", low_start, low_end, step)
hi_indices = g.op("Range", size, hi_end, step)
low_size = symbolic_helper._size_helper(
g, low_indices, g.op("Constant", value_t=torch.tensor(0)))
hi_size = symbolic_helper._size_helper(
g, hi_indices, g.op("Constant", value_t=torch.tensor(0)))
ndim = symbolic_helper._get_tensor_rank(input)
assert ndim is not None
perm = list(range(0, ndim))
perm.append(perm.pop(dimension))
unsqueeze_list = []
loop_condition = g.op("Constant", value_t=torch.tensor(1))
loop_condition = g.op("Cast", loop_condition, to_i=9)
loop_len = g.op("Min", low_size, hi_size)
loop = g.op("Loop", loop_len, loop_condition)
loop_block = utils._add_block(loop.node())
block_input_iter = utils._add_input_to_block(loop_block)
cond = utils._add_input_to_block(loop_block)
starts = loop_block.op("Gather", low_indices, block_input_iter)
ends = loop_block.op("Gather", hi_indices, block_input_iter)
axes = loop_block.op("Constant", value_t=torch.tensor([2]))
starts = symbolic_helper._unsqueeze_helper(loop_block, starts, [0])
ends = symbolic_helper._unsqueeze_helper(loop_block, ends, [0])
stack = loop_block.op("Slice", input, starts, ends, axes)
unsqueeze = symbolic_helper._unsqueeze_helper(
loop_block, loop_block.op("Transpose", stack, perm_i=perm),
[dimension])
unsqueeze_list.append(unsqueeze)
concat = loop_block.op("Concat", *unsqueeze_list, axis_i=0)
cond_out = loop_block.op("Cast", loop_condition, to_i=9)
utils._add_output_to_block(loop_block, cond_out)
utils._add_output_to_block(loop_block, concat)
loop_output = loop.node().output()
perm = [0, 1, 2, 3, 4]
perm[0], perm[dimension + 1] = perm[dimension + 1], perm[0]
transpose = g.op("Transpose", loop_output, perm_i=perm)
squeeze = symbolic_helper._squeeze_helper(g, transpose, [0])
return squeeze
else:
return symbolic_helper._unimplemented("Unfold",
"input size not accessible")
