def embedding_bag(g,
embedding_matrix,
indices,
offsets,
scale_grad_by_freq,
mode,
sparse,
per_sample_weights,
include_last_offset,
padding_idx):
if scale_grad_by_freq and sym_help._training_mode:
return sym_help._onnx_unsupported("embedding_bag with scale_grad_by_freq for training mode")
if padding_idx is not None and padding_idx >= 0:
raise RuntimeError("embedding_bag with padding_idx")
from torch.onnx.symbolic_opset9 import select
import warnings
warnings.warn("Export of embedding_bag with dynamic input/offsets shape is not supported in opset 10. "
"Please use opset 11 or higher to export model for dynamic input shape.'")
offsets_dim_0 = sym_help._get_tensor_dim_size(offsets, 0)
if offsets_dim_0 is not None:
if include_last_offset:
offset_len = offsets_dim_0 - 1
offsets_extended = offsets
else:
offset_len = offsets_dim_0
offsets_extended = [offsets, g.op("Constant", value_t=torch.tensor([maxsize]))]
offsets_extended = g.op("Concat", *offsets_extended, axis_i=0)
list_ = []
for i in range(offset_len):
start_ = sym_help._unsqueeze_helper(g, select(g, offsets_extended, torch.tensor(0), torch.tensor(i)), [0])
end_ = sym_help._unsqueeze_helper(g, select(g, offsets_extended, torch.tensor(0), torch.tensor(i + 1)), [0])
axes_ = g.op("Constant", value_t=torch.tensor([0]))
indices_row = g.op("Slice", indices, start_, end_, axes_)
embeddings = g.op("Gather", embedding_matrix, indices_row)
if not sym_help._is_none(per_sample_weights):
per_sample_weights_row = g.op("Slice", per_sample_weights, start_, end_, axes_)
per_sample_weights_row = sym_help._unsqueeze_helper(g, per_sample_weights_row, [1])
embeddings = g.op("Mul", embeddings, per_sample_weights_row)
if mode == 0:
embeddings = sym_help._reducesum_helper(g, embeddings, axes_i=[0], keepdims_i=0)
elif mode == 1:
embeddings = g.op("ReduceMean", embeddings, axes_i=[0], keepdims_i=0)
else:
embeddings = g.op("ReduceMax", embeddings, axes_i=[0], keepdims_i=0)
embeddings = sym_help._unsqueeze_helper(g, embeddings, [0])
list_.append(embeddings)
output = g.op("Concat", *list_, axis_i=0)
# aten::embedding_bag returns a tuple of 4 elements: output, offset2bag, bag_size, max_indices.
# But the last three outputs are not used in torch.nn.EmbeddingBag or torch.nn.functional.embedding_bag.
return output, None, None, None
else:
return sym_help._onnx_unsupported("embedding_bag with unknown shape of offsets for opset 10 is not supported. "
"please use opset 11 or higher.")
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 embedding_bag(g, embedding_matrix, indices, offsets, scale_grad_by_freq,
mode, sparse, per_sample_weights, include_last_offset,
padding_idx):
if scale_grad_by_freq and sym_help._training_mode:
return sym_help._onnx_unsupported(
'embedding_bag with scale_grad_by_freq for training mode')
if padding_idx is not None and padding_idx >= 0:
raise RuntimeError('embedding_bag with padding_idx')
loop_condition = g.op("Constant", value_t=torch.tensor(1))
loop_condition = g.op("Cast", loop_condition, to_i=9)
zero = g.op("Constant", value_t=torch.tensor([0]))
indices_len = sym_help._unsqueeze_helper(
g,
sym_help._size_helper(g, indices,
g.op("Constant", value_t=torch.tensor(0))), [0])
if not include_last_offset:
offsets = [offsets, indices_len]
offsets = g.op("Concat", *offsets, axis_i=0)
# Offsets holds the starting index position of each bag. So we create a list of the indices slices (determined by
# offsets) and gather those indices in indices_row. Then we use this subset of indices to gather from embeddings.
# The embeddings output is a loop scan output, so we can avoid creating a sequence and inserting elements in.
offsets_starts = sym_help._slice_helper(g,
offsets,
axes=[0],
starts=[0],
ends=[maxsize],
steps=[1])
offsets_ends = sym_help._slice_helper(g,
offsets,
axes=[0],
starts=[1],
ends=[maxsize],
steps=[1])
loop_len = sym_help._size_helper(g, offsets_ends,
g.op("Constant", value_t=torch.tensor(0)))
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)
indices_start = loop_block.op("Gather",
offsets_starts,
block_input_iter,
axis_i=0)
indices_end = loop_block.op("Gather",
offsets_ends,
block_input_iter,
axis_i=0)
indices_start = sym_help._unsqueeze_helper(loop_block, indices_start, [0])
indices_end = sym_help._unsqueeze_helper(loop_block, indices_end, [0])
indices_row = loop_block.op("Slice", indices, indices_start, indices_end,
zero)
embeddings = loop_block.op("Gather",
embedding_matrix,
indices_row,
axis_i=0)
if not sym_help._is_none(per_sample_weights):
per_sample_weights_row = loop_block.op("Slice", per_sample_weights,
indices_start, indices_end,
zero)
per_sample_weights_row = sym_help._unsqueeze_helper(
loop_block, per_sample_weights_row, [1])
embeddings = loop_block.op("Mul", embeddings, per_sample_weights_row)
if mode == 0:
embeddings = sym_help._reducesum_helper(loop_block,
embeddings,
axes_i=[0],
keepdims_i=0)
elif mode == 1:
embeddings = loop_block.op("ReduceMean",
embeddings,
axes_i=[0],
keepdims_i=0)
else:
embeddings = loop_block.op("ReduceMax",
embeddings,
axes_i=[0],
keepdims_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, embeddings)
# aten::embedding_bag returns a tuple of 4 elements: output, offset2bag, bag_size, max_indices.
# But the last three outputs are not used in torch.nn.EmbeddingBag or torch.nn.functional.embedding_bag.
return loop.node().output(), None, None, None
def bad_clamp(g, self, min, max):
return symbolic_helper._onnx_unsupported("Bad boy!")
def bad_clamp(g, self, min, max):
return _onnx_unsupported("Bad boy!")
def embedding_bag(g, embedding_matrix, indices, offsets, scale_grad_by_freq,
mode, sparse, per_sample_weights, include_last_offset):
if scale_grad_by_freq and sym_help._training_mode:
return sym_help._onnx_unsupported(
'embedding_bag with scale_grad_by_freq for training mode')
from torch.onnx.symbolic_opset9 import size, div, select
# Check if initial indices was 2D. In functional.py:
# offsets is set to torch.arange(0, indices.numel(), indices.size(1))
# Then indices is reshaped to 1D: indices.reshape(-1)
if len(list(indices.node().inputs())) > 0 and indices.node().inputs().__next__().type().sizes() is not None \
and len(indices.node().inputs().__next__().type().sizes()) == 2:
# Assert include_last_offset is False
assert not include_last_offset
embeddings = g.op("Gather", embedding_matrix, indices)
dim_0 = size(g, offsets, g.op("Constant",
value_t=torch.LongTensor([0])))
dim_1 = div(
g, size(g, indices, g.op("Constant",
value_t=torch.LongTensor([0]))), dim_0)
dim_2 = g.op("Constant", value_t=torch.LongTensor([-1]))
shape = [dim_0, dim_1, dim_2]
shape = g.op("Concat", *shape, axis_i=0)
if not sym_help._is_none(per_sample_weights):
per_sample_weights = g.op("Unsqueeze",
per_sample_weights,
axes_i=[1])
embeddings = g.op("Mul", embeddings, per_sample_weights)
embeddings = g.op("Reshape", embeddings, shape)
if mode == 0:
embeddings = g.op("ReduceSum",
embeddings,
axes_i=[1],
keepdims_i=0)
elif mode == 1:
embeddings = g.op("ReduceMean",
embeddings,
axes_i=[1],
keepdims_i=0)
else:
embeddings = g.op("ReduceMax",
embeddings,
axes_i=[1],
keepdims_i=0)
# aten::embedding_bag returns a tuple of 4 elements: output, offset2bag, bag_size, max_indices.
# But the last three outputs are not used in torch.nn.EmbeddingBag or torch.nn.functional.embedding_bag.
return embeddings, None, None, None
elif offsets.type().sizes() is not None:
if include_last_offset:
offset_len = offsets.type().sizes()[0] - 1
offsets_extended = offsets
else:
offset_len = offsets.type().sizes()[0]
offsets_extended = [
offsets,
g.op("Constant", value_t=torch.tensor([maxsize]))
]
offsets_extended = g.op("Concat", *offsets_extended, axis_i=0)
list_ = []
for i in range(offset_len):
start_ = g.op("Unsqueeze",
select(g, offsets_extended, torch.tensor(0),
torch.tensor(i)),
axes_i=[0])
end_ = g.op("Unsqueeze",
select(g, offsets_extended, torch.tensor(0),
torch.tensor(i + 1)),
axes_i=[0])
axes_ = g.op("Constant", value_t=torch.tensor([0]))
indices_row = g.op("Slice", indices, start_, end_, axes_)
embeddings = g.op("Gather", embedding_matrix, indices_row)
if not sym_help._is_none(per_sample_weights):
per_sample_weights_row = g.op("Slice", per_sample_weights,
start_, end_, axes_)
per_sample_weights_row = g.op("Unsqueeze",
per_sample_weights_row,
axes_i=[1])
embeddings = g.op("Mul", embeddings, per_sample_weights_row)
if mode == 0:
embeddings = g.op("ReduceSum",
embeddings,
axes_i=[0],
keepdims_i=0)
elif mode == 1:
embeddings = g.op("ReduceMean",
embeddings,
axes_i=[0],
keepdims_i=0)
else:
embeddings = g.op("ReduceMax",
embeddings,
axes_i=[0],
keepdims_i=0)
embeddings = g.op("Unsqueeze", embeddings, axes_i=[0])
list_.append(embeddings)
output = g.op("Concat", *list_, axis_i=0)
# aten::embedding_bag returns a tuple of 4 elements: output, offset2bag, bag_size, max_indices.
# But the last three outputs are not used in torch.nn.EmbeddingBag or torch.nn.functional.embedding_bag.
return output, None, None, None
else:
return sym_help._onnx_unsupported(
'embedding_bag with unknown shape of indices')
