def margin_ranking_loss(
input1: TensorLikeType,
input2: TensorLikeType,
target: TensorLikeType,
margin: float = 0.0,
reduction: str = "mean",
) -> TensorLikeType:
# Formula of loss (implementation gets confusing with all the refs.foo)
# loss_without_reduction = max(0, −target * (input1 − input2) + margin)
if input1.ndim != input2.ndim or input1.ndim != target.ndim:
raise RuntimeError(
(
"margin_ranking_loss : All input tensors should have same dimension but got sizes: "
"input1: {}, input2: {}, target: {} ".format(
input1.shape, input2.shape, target.shape
)
)
)
_check_reduction_value(reduction)
neg_target = refs.neg(target)
input_diff = refs.sub(input1, input2)
mul_target_input = refs.mul(neg_target, input_diff)
add_margin = refs.add(mul_target_input, margin)
loss = refs.maximum(add_margin, 0)
return _apply_loss_reduction(loss, reduction)
def hinge_embedding_loss(
input: TensorLikeType,
target: TensorLikeType,
margin: float = 1.0,
reduction: str = "mean",
) -> TensorLikeType:
# Formula of loss (implementation gets confusing with all the refs.foo)
# loss_without_reduction = input if y == 1
# = max(0, margin - input) if y == -1
_check_reduction_value(reduction)
margin_clamp = refs.maximum(refs.sub(margin, input), 0)
output_margin = refs.where(refs.ne(target, 1), margin_clamp, 0)
output_self = refs.where(refs.ne(target, -1), input, 0)
loss = refs.add(output_margin, output_self)
return _apply_loss_reduction(loss, reduction)
fusion1.print_ir()
# Execute Fusion
input1 = torch.randn(3, device='cuda')
input2 = torch.randn(2, 3, 4, device='cuda')
# Kernel compilation should be cached for the 2nd iteration
# with input tensors of the same shape
for _ in range(5) :
o = fusion1.execute([input1, input2])[0]
assert(o.shape == torch.Size([2, 3, 4]))
# Reference in prim torch
ref_o = refs.add(prims.broadcast_in_dim(input1, [2, 3, 4], [1]), input2)
assert(ref_o.allclose(o))
assert(ref_o.shape == o.shape)
fusion2 = Fusion()
input1 = torch.randn(1, 1, 4, device='cuda')
input2 = torch.randn(2, 3, 4, device='cuda')
with FusionDefinition(fusion2) as fd :
t0 = fd.define_tensor(sizes=input1.size(), strides=input1.stride())
t1 = fd.define_tensor(sizes=input2.size(), strides=input2.stride())
t0_b = fd.ops.broadcast_in_dim(t0, [2, 3, 4], [0, 1, 2])
t2 = fd.ops.add(t0_b, t1)