chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,14 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import importlib
|
||||
import os
|
||||
|
||||
|
||||
# automatically import any Python files in the criterions/ directory
|
||||
for file in os.listdir(os.path.dirname(__file__)):
|
||||
if file.endswith(".py") and not file.startswith("_"):
|
||||
module = file[: file.find(".py")]
|
||||
importlib.import_module("examples.simultaneous_translation.utils." + module)
|
||||
@@ -0,0 +1,149 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def exclusive_cumprod(tensor, dim: int, eps: float = 1e-10):
|
||||
"""
|
||||
Implementing exclusive cumprod.
|
||||
There is cumprod in pytorch, however there is no exclusive mode.
|
||||
cumprod(x) = [x1, x1x2, x2x3x4, ..., prod_{i=1}^n x_i]
|
||||
exclusive means cumprod(x) = [1, x1, x1x2, x1x2x3, ..., prod_{i=1}^{n-1} x_i]
|
||||
"""
|
||||
tensor_size = list(tensor.size())
|
||||
tensor_size[dim] = 1
|
||||
return_tensor = safe_cumprod(
|
||||
torch.cat([torch.ones(tensor_size).type_as(tensor), tensor], dim=dim),
|
||||
dim=dim,
|
||||
eps=eps,
|
||||
)
|
||||
|
||||
if dim == 0:
|
||||
return return_tensor[:-1]
|
||||
elif dim == 1:
|
||||
return return_tensor[:, :-1]
|
||||
elif dim == 2:
|
||||
return return_tensor[:, :, :-1]
|
||||
else:
|
||||
raise RuntimeError("Cumprod on dimension 3 and more is not implemented")
|
||||
|
||||
|
||||
def safe_cumprod(tensor, dim: int, eps: float = 1e-10):
|
||||
"""
|
||||
An implementation of cumprod to prevent precision issue.
|
||||
cumprod(x)
|
||||
= [x1, x1x2, x1x2x3, ....]
|
||||
= [exp(log(x1)), exp(log(x1) + log(x2)), exp(log(x1) + log(x2) + log(x3)), ...]
|
||||
= exp(cumsum(log(x)))
|
||||
"""
|
||||
|
||||
if (tensor + eps < 0).any().item():
|
||||
raise RuntimeError(
|
||||
"Safe cumprod can only take non-negative tensors as input."
|
||||
"Consider use torch.cumprod if you want to calculate negative values."
|
||||
)
|
||||
|
||||
log_tensor = torch.log(tensor + eps)
|
||||
cumsum_log_tensor = torch.cumsum(log_tensor, dim)
|
||||
exp_cumsum_log_tensor = torch.exp(cumsum_log_tensor)
|
||||
return exp_cumsum_log_tensor
|
||||
|
||||
|
||||
def lengths_to_mask(lengths, max_len: int, dim: int = 0, negative_mask: bool = False):
|
||||
"""
|
||||
Convert a tensor of lengths to mask
|
||||
For example, lengths = [[2, 3, 4]], max_len = 5
|
||||
mask =
|
||||
[[1, 1, 1],
|
||||
[1, 1, 1],
|
||||
[0, 1, 1],
|
||||
[0, 0, 1],
|
||||
[0, 0, 0]]
|
||||
"""
|
||||
assert len(lengths.size()) <= 2
|
||||
if len(lengths) == 2:
|
||||
if dim == 1:
|
||||
lengths = lengths.t()
|
||||
lengths = lengths
|
||||
else:
|
||||
lengths = lengths.unsqueeze(1)
|
||||
|
||||
# lengths : batch_size, 1
|
||||
lengths = lengths.view(-1, 1)
|
||||
|
||||
batch_size = lengths.size(0)
|
||||
# batch_size, max_len
|
||||
mask = torch.arange(max_len).expand(batch_size, max_len).type_as(lengths) < lengths
|
||||
|
||||
if negative_mask:
|
||||
mask = ~mask
|
||||
|
||||
if dim == 0:
|
||||
# max_len, batch_size
|
||||
mask = mask.t()
|
||||
|
||||
return mask
|
||||
|
||||
|
||||
def moving_sum(x, start_idx: int, end_idx: int):
|
||||
"""
|
||||
From MONOTONIC CHUNKWISE ATTENTION
|
||||
https://arxiv.org/pdf/1712.05382.pdf
|
||||
Equation (18)
|
||||
|
||||
x = [x_1, x_2, ..., x_N]
|
||||
MovingSum(x, start_idx, end_idx)_n = Sigma_{m=n−(start_idx−1)}^{n+end_idx-1} x_m
|
||||
for n in {1, 2, 3, ..., N}
|
||||
|
||||
x : src_len, batch_size
|
||||
start_idx : start idx
|
||||
end_idx : end idx
|
||||
|
||||
Example
|
||||
src_len = 5
|
||||
batch_size = 3
|
||||
x =
|
||||
[[ 0, 5, 10],
|
||||
[ 1, 6, 11],
|
||||
[ 2, 7, 12],
|
||||
[ 3, 8, 13],
|
||||
[ 4, 9, 14]]
|
||||
|
||||
MovingSum(x, 3, 1) =
|
||||
[[ 0, 5, 10],
|
||||
[ 1, 11, 21],
|
||||
[ 3, 18, 33],
|
||||
[ 6, 21, 36],
|
||||
[ 9, 24, 39]]
|
||||
|
||||
MovingSum(x, 1, 3) =
|
||||
[[ 3, 18, 33],
|
||||
[ 6, 21, 36],
|
||||
[ 9, 24, 39],
|
||||
[ 7, 17, 27],
|
||||
[ 4, 9, 14]]
|
||||
"""
|
||||
assert start_idx > 0 and end_idx > 0
|
||||
assert len(x.size()) == 2
|
||||
src_len, batch_size = x.size()
|
||||
# batch_size, 1, src_len
|
||||
x = x.t().unsqueeze(1)
|
||||
# batch_size, 1, src_len
|
||||
moving_sum_weight = x.new_ones([1, 1, end_idx + start_idx - 1])
|
||||
|
||||
moving_sum = (
|
||||
torch.nn.functional.conv1d(
|
||||
x, moving_sum_weight, padding=start_idx + end_idx - 1
|
||||
)
|
||||
.squeeze(1)
|
||||
.t()
|
||||
)
|
||||
moving_sum = moving_sum[end_idx:-start_idx]
|
||||
|
||||
assert src_len == moving_sum.size(0)
|
||||
assert batch_size == moving_sum.size(1)
|
||||
|
||||
return moving_sum
|
||||
@@ -0,0 +1,451 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class LatencyMetric(object):
|
||||
@staticmethod
|
||||
def length_from_padding_mask(padding_mask, batch_first: bool = False):
|
||||
dim = 1 if batch_first else 0
|
||||
return padding_mask.size(dim) - padding_mask.sum(dim=dim, keepdim=True)
|
||||
|
||||
def prepare_latency_metric(
|
||||
self,
|
||||
delays,
|
||||
src_lens,
|
||||
target_padding_mask=None,
|
||||
batch_first: bool = False,
|
||||
start_from_zero: bool = True,
|
||||
):
|
||||
assert len(delays.size()) == 2
|
||||
assert len(src_lens.size()) == 2
|
||||
|
||||
if start_from_zero:
|
||||
delays = delays + 1
|
||||
|
||||
if batch_first:
|
||||
# convert to batch_last
|
||||
delays = delays.t()
|
||||
src_lens = src_lens.t()
|
||||
tgt_len, bsz = delays.size()
|
||||
_, bsz_1 = src_lens.size()
|
||||
|
||||
if target_padding_mask is not None:
|
||||
target_padding_mask = target_padding_mask.t()
|
||||
tgt_len_1, bsz_2 = target_padding_mask.size()
|
||||
assert tgt_len == tgt_len_1
|
||||
assert bsz == bsz_2
|
||||
|
||||
assert bsz == bsz_1
|
||||
|
||||
if target_padding_mask is None:
|
||||
tgt_lens = tgt_len * delays.new_ones([1, bsz]).float()
|
||||
else:
|
||||
# 1, batch_size
|
||||
tgt_lens = self.length_from_padding_mask(target_padding_mask, False).float()
|
||||
delays = delays.masked_fill(target_padding_mask, 0)
|
||||
|
||||
return delays, src_lens, tgt_lens, target_padding_mask
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
delays,
|
||||
src_lens,
|
||||
target_padding_mask=None,
|
||||
batch_first: bool = False,
|
||||
start_from_zero: bool = True,
|
||||
):
|
||||
delays, src_lens, tgt_lens, target_padding_mask = self.prepare_latency_metric(
|
||||
delays, src_lens, target_padding_mask, batch_first, start_from_zero
|
||||
)
|
||||
return self.cal_metric(delays, src_lens, tgt_lens, target_padding_mask)
|
||||
|
||||
@staticmethod
|
||||
def cal_metric(delays, src_lens, tgt_lens, target_padding_mask):
|
||||
"""
|
||||
Expected sizes:
|
||||
delays: tgt_len, batch_size
|
||||
src_lens: 1, batch_size
|
||||
target_padding_mask: tgt_len, batch_size
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class AverageProportion(LatencyMetric):
|
||||
"""
|
||||
Function to calculate Average Proportion from
|
||||
Can neural machine translation do simultaneous translation?
|
||||
(https://arxiv.org/abs/1606.02012)
|
||||
|
||||
Delays are monotonic steps, range from 1 to src_len.
|
||||
Give src x tgt y, AP is calculated as:
|
||||
|
||||
AP = 1 / (|x||y]) sum_i^|Y| deleys_i
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def cal_metric(delays, src_lens, tgt_lens, target_padding_mask):
|
||||
if target_padding_mask is not None:
|
||||
AP = torch.sum(
|
||||
delays.masked_fill(target_padding_mask, 0), dim=0, keepdim=True
|
||||
)
|
||||
else:
|
||||
AP = torch.sum(delays, dim=0, keepdim=True)
|
||||
|
||||
AP = AP / (src_lens * tgt_lens)
|
||||
return AP
|
||||
|
||||
|
||||
class AverageLagging(LatencyMetric):
|
||||
"""
|
||||
Function to calculate Average Lagging from
|
||||
STACL: Simultaneous Translation with Implicit Anticipation
|
||||
and Controllable Latency using Prefix-to-Prefix Framework
|
||||
(https://arxiv.org/abs/1810.08398)
|
||||
|
||||
Delays are monotonic steps, range from 1 to src_len.
|
||||
Give src x tgt y, AP is calculated as:
|
||||
|
||||
AL = 1 / tau sum_i^tau delays_i - (i - 1) / gamma
|
||||
|
||||
Where
|
||||
gamma = |y| / |x|
|
||||
tau = argmin_i(delays_i = |x|)
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def cal_metric(delays, src_lens, tgt_lens, target_padding_mask):
|
||||
# tau = argmin_i(delays_i = |x|)
|
||||
tgt_len, bsz = delays.size()
|
||||
lagging_padding_mask = delays >= src_lens
|
||||
lagging_padding_mask = torch.nn.functional.pad(
|
||||
lagging_padding_mask.t(), (1, 0)
|
||||
).t()[:-1, :]
|
||||
gamma = tgt_lens / src_lens
|
||||
lagging = (
|
||||
delays
|
||||
- torch.arange(delays.size(0))
|
||||
.unsqueeze(1)
|
||||
.type_as(delays)
|
||||
.expand_as(delays)
|
||||
/ gamma
|
||||
)
|
||||
lagging.masked_fill_(lagging_padding_mask, 0)
|
||||
tau = (1 - lagging_padding_mask.type_as(lagging)).sum(dim=0, keepdim=True)
|
||||
AL = lagging.sum(dim=0, keepdim=True) / tau
|
||||
|
||||
return AL
|
||||
|
||||
|
||||
class DifferentiableAverageLagging(LatencyMetric):
|
||||
"""
|
||||
Function to calculate Differentiable Average Lagging from
|
||||
Monotonic Infinite Lookback Attention for Simultaneous Machine Translation
|
||||
(https://arxiv.org/abs/1906.05218)
|
||||
|
||||
Delays are monotonic steps, range from 0 to src_len-1.
|
||||
(In the original paper thery are from 1 to src_len)
|
||||
Give src x tgt y, AP is calculated as:
|
||||
|
||||
DAL = 1 / |Y| sum_i^|Y| delays'_i - (i - 1) / gamma
|
||||
|
||||
Where
|
||||
delays'_i =
|
||||
1. delays_i if i == 1
|
||||
2. max(delays_i, delays'_{i-1} + 1 / gamma)
|
||||
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def cal_metric(delays, src_lens, tgt_lens, target_padding_mask):
|
||||
tgt_len, bsz = delays.size()
|
||||
|
||||
gamma = tgt_lens / src_lens
|
||||
new_delays = torch.zeros_like(delays)
|
||||
|
||||
for i in range(delays.size(0)):
|
||||
if i == 0:
|
||||
new_delays[i] = delays[i]
|
||||
else:
|
||||
new_delays[i] = torch.cat(
|
||||
[
|
||||
new_delays[i - 1].unsqueeze(0) + 1 / gamma,
|
||||
delays[i].unsqueeze(0),
|
||||
],
|
||||
dim=0,
|
||||
).max(dim=0)[0]
|
||||
|
||||
DAL = (
|
||||
new_delays
|
||||
- torch.arange(delays.size(0))
|
||||
.unsqueeze(1)
|
||||
.type_as(delays)
|
||||
.expand_as(delays)
|
||||
/ gamma
|
||||
)
|
||||
if target_padding_mask is not None:
|
||||
DAL = DAL.masked_fill(target_padding_mask, 0)
|
||||
|
||||
DAL = DAL.sum(dim=0, keepdim=True) / tgt_lens
|
||||
|
||||
return DAL
|
||||
|
||||
|
||||
class LatencyMetricVariance(LatencyMetric):
|
||||
def prepare_latency_metric(
|
||||
self,
|
||||
delays,
|
||||
src_lens,
|
||||
target_padding_mask=None,
|
||||
batch_first: bool = True,
|
||||
start_from_zero: bool = True,
|
||||
):
|
||||
assert batch_first
|
||||
assert len(delays.size()) == 3
|
||||
assert len(src_lens.size()) == 2
|
||||
|
||||
if start_from_zero:
|
||||
delays = delays + 1
|
||||
|
||||
# convert to batch_last
|
||||
bsz, num_heads_x_layers, tgt_len = delays.size()
|
||||
bsz_1, _ = src_lens.size()
|
||||
assert bsz == bsz_1
|
||||
|
||||
if target_padding_mask is not None:
|
||||
bsz_2, tgt_len_1 = target_padding_mask.size()
|
||||
assert tgt_len == tgt_len_1
|
||||
assert bsz == bsz_2
|
||||
|
||||
if target_padding_mask is None:
|
||||
tgt_lens = tgt_len * delays.new_ones([bsz, tgt_len]).float()
|
||||
else:
|
||||
# batch_size, 1
|
||||
tgt_lens = self.length_from_padding_mask(target_padding_mask, True).float()
|
||||
delays = delays.masked_fill(target_padding_mask.unsqueeze(1), 0)
|
||||
|
||||
return delays, src_lens, tgt_lens, target_padding_mask
|
||||
|
||||
|
||||
class VarianceDelay(LatencyMetricVariance):
|
||||
@staticmethod
|
||||
def cal_metric(delays, src_lens, tgt_lens, target_padding_mask):
|
||||
"""
|
||||
delays : bsz, num_heads_x_layers, tgt_len
|
||||
src_lens : bsz, 1
|
||||
target_lens : bsz, 1
|
||||
target_padding_mask: bsz, tgt_len or None
|
||||
"""
|
||||
if delays.size(1) == 1:
|
||||
return delays.new_zeros([1])
|
||||
|
||||
variance_delays = delays.var(dim=1)
|
||||
|
||||
if target_padding_mask is not None:
|
||||
variance_delays.masked_fill_(target_padding_mask, 0)
|
||||
|
||||
return variance_delays.sum(dim=1, keepdim=True) / tgt_lens
|
||||
|
||||
|
||||
class LatencyInference(object):
|
||||
def __init__(self, start_from_zero=True):
|
||||
self.metric_calculator = {
|
||||
"differentiable_average_lagging": DifferentiableAverageLagging(),
|
||||
"average_lagging": AverageLagging(),
|
||||
"average_proportion": AverageProportion(),
|
||||
}
|
||||
|
||||
self.start_from_zero = start_from_zero
|
||||
|
||||
def __call__(self, monotonic_step, src_lens):
|
||||
"""
|
||||
monotonic_step range from 0 to src_len. src_len means eos
|
||||
delays: bsz, tgt_len
|
||||
src_lens: bsz, 1
|
||||
"""
|
||||
if not self.start_from_zero:
|
||||
monotonic_step -= 1
|
||||
|
||||
src_lens = src_lens
|
||||
|
||||
delays = monotonic_step.view(
|
||||
monotonic_step.size(0), -1, monotonic_step.size(-1)
|
||||
).max(dim=1)[0]
|
||||
|
||||
delays = delays.masked_fill(delays >= src_lens, 0) + (src_lens - 1).expand_as(
|
||||
delays
|
||||
).masked_fill(delays < src_lens, 0)
|
||||
return_dict = {}
|
||||
for key, func in self.metric_calculator.items():
|
||||
return_dict[key] = func(
|
||||
delays.float(),
|
||||
src_lens.float(),
|
||||
target_padding_mask=None,
|
||||
batch_first=True,
|
||||
start_from_zero=True,
|
||||
).t()
|
||||
|
||||
return return_dict
|
||||
|
||||
|
||||
class LatencyTraining(object):
|
||||
def __init__(
|
||||
self,
|
||||
avg_weight,
|
||||
var_weight,
|
||||
avg_type,
|
||||
var_type,
|
||||
stay_on_last_token,
|
||||
average_method,
|
||||
):
|
||||
self.avg_weight = avg_weight
|
||||
self.var_weight = var_weight
|
||||
self.avg_type = avg_type
|
||||
self.var_type = var_type
|
||||
self.stay_on_last_token = stay_on_last_token
|
||||
self.average_method = average_method
|
||||
|
||||
self.metric_calculator = {
|
||||
"differentiable_average_lagging": DifferentiableAverageLagging(),
|
||||
"average_lagging": AverageLagging(),
|
||||
"average_proportion": AverageProportion(),
|
||||
}
|
||||
|
||||
self.variance_calculator = {
|
||||
"variance_delay": VarianceDelay(),
|
||||
}
|
||||
|
||||
def expected_delays_from_attention(
|
||||
self, attention, source_padding_mask=None, target_padding_mask=None
|
||||
):
|
||||
if type(attention) == list:
|
||||
# bsz, num_heads, tgt_len, src_len
|
||||
bsz, num_heads, tgt_len, src_len = attention[0].size()
|
||||
attention = torch.cat(attention, dim=1)
|
||||
bsz, num_heads_x_layers, tgt_len, src_len = attention.size()
|
||||
# bsz * num_heads * num_layers, tgt_len, src_len
|
||||
attention = attention.view(-1, tgt_len, src_len)
|
||||
else:
|
||||
# bsz * num_heads * num_layers, tgt_len, src_len
|
||||
bsz, tgt_len, src_len = attention.size()
|
||||
num_heads_x_layers = 1
|
||||
attention = attention.view(-1, tgt_len, src_len)
|
||||
|
||||
if not self.stay_on_last_token:
|
||||
residual_attention = 1 - attention[:, :, :-1].sum(dim=2, keepdim=True)
|
||||
attention = torch.cat([attention[:, :, :-1], residual_attention], dim=2)
|
||||
|
||||
# bsz * num_heads_x_num_layers, tgt_len, src_len for MMA
|
||||
steps = (
|
||||
torch.arange(1, 1 + src_len)
|
||||
.unsqueeze(0)
|
||||
.unsqueeze(1)
|
||||
.expand_as(attention)
|
||||
.type_as(attention)
|
||||
)
|
||||
|
||||
if source_padding_mask is not None:
|
||||
src_offset = (
|
||||
source_padding_mask.type_as(attention)
|
||||
.sum(dim=1, keepdim=True)
|
||||
.expand(bsz, num_heads_x_layers)
|
||||
.contiguous()
|
||||
.view(-1, 1)
|
||||
)
|
||||
src_lens = src_len - src_offset
|
||||
if source_padding_mask[:, 0].any():
|
||||
# Pad left
|
||||
src_offset = src_offset.view(-1, 1, 1)
|
||||
steps = steps - src_offset
|
||||
steps = steps.masked_fill(steps <= 0, 0)
|
||||
else:
|
||||
src_lens = attention.new_ones([bsz, num_heads_x_layers]) * src_len
|
||||
src_lens = src_lens.view(-1, 1)
|
||||
|
||||
# bsz * num_heads_num_layers, tgt_len, src_len
|
||||
expected_delays = (
|
||||
(steps * attention).sum(dim=2).view(bsz, num_heads_x_layers, tgt_len)
|
||||
)
|
||||
|
||||
if target_padding_mask is not None:
|
||||
expected_delays.masked_fill_(target_padding_mask.unsqueeze(1), 0)
|
||||
|
||||
return expected_delays, src_lens
|
||||
|
||||
def avg_loss(self, expected_delays, src_lens, target_padding_mask):
|
||||
|
||||
bsz, num_heads_x_layers, tgt_len = expected_delays.size()
|
||||
target_padding_mask = (
|
||||
target_padding_mask.unsqueeze(1)
|
||||
.expand_as(expected_delays)
|
||||
.contiguous()
|
||||
.view(-1, tgt_len)
|
||||
)
|
||||
|
||||
if self.average_method == "average":
|
||||
# bsz * tgt_len
|
||||
expected_delays = expected_delays.mean(dim=1)
|
||||
elif self.average_method == "weighted_average":
|
||||
weights = torch.nn.functional.softmax(expected_delays, dim=1)
|
||||
expected_delays = torch.sum(expected_delays * weights, dim=1)
|
||||
elif self.average_method == "max":
|
||||
# bsz * num_heads_x_num_layers, tgt_len
|
||||
expected_delays = expected_delays.max(dim=1)[0]
|
||||
else:
|
||||
raise RuntimeError(f"{self.average_method} is not supported")
|
||||
|
||||
src_lens = src_lens.view(bsz, -1)[:, :1]
|
||||
target_padding_mask = target_padding_mask.view(bsz, -1, tgt_len)[:, 0]
|
||||
|
||||
if self.avg_weight > 0.0:
|
||||
if self.avg_type in self.metric_calculator:
|
||||
average_delays = self.metric_calculator[self.avg_type](
|
||||
expected_delays,
|
||||
src_lens,
|
||||
target_padding_mask,
|
||||
batch_first=True,
|
||||
start_from_zero=False,
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(f"{self.avg_type} is not supported.")
|
||||
|
||||
# bsz * num_heads_x_num_layers, 1
|
||||
return self.avg_weight * average_delays.sum()
|
||||
else:
|
||||
return 0.0
|
||||
|
||||
def var_loss(self, expected_delays, src_lens, target_padding_mask):
|
||||
src_lens = src_lens.view(expected_delays.size(0), expected_delays.size(1))[
|
||||
:, :1
|
||||
]
|
||||
if self.var_weight > 0.0:
|
||||
if self.var_type in self.variance_calculator:
|
||||
variance_delays = self.variance_calculator[self.var_type](
|
||||
expected_delays,
|
||||
src_lens,
|
||||
target_padding_mask,
|
||||
batch_first=True,
|
||||
start_from_zero=False,
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(f"{self.var_type} is not supported.")
|
||||
|
||||
return self.var_weight * variance_delays.sum()
|
||||
else:
|
||||
return 0.0
|
||||
|
||||
def loss(self, attention, source_padding_mask=None, target_padding_mask=None):
|
||||
expected_delays, src_lens = self.expected_delays_from_attention(
|
||||
attention, source_padding_mask, target_padding_mask
|
||||
)
|
||||
|
||||
latency_loss = 0
|
||||
|
||||
latency_loss += self.avg_loss(expected_delays, src_lens, target_padding_mask)
|
||||
|
||||
latency_loss += self.var_loss(expected_delays, src_lens, target_padding_mask)
|
||||
|
||||
return latency_loss
|
||||
Reference in New Issue
Block a user