Skip to content

Commit 84add25

Browse files
flexthinkflexthink
committed
Add a guided attention loss
1 parent 0d3ee04 commit 84add25

2 files changed

Lines changed: 236 additions & 0 deletions

File tree

‎speechbrain/nnet/loss/guidedattn_loss.py‎

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
"""The Guided Attention Loss implementation
2+
3+
This loss can be used to speed up the training of
4+
models in which the correspondence between inputs and
5+
outputs is roughly linear, and the attention alignments
6+
are expected to be approximately diagonal, such as Grapheme-to-Phoneme
7+
and Text-to-Speech
8+
9+
Authors
10+
* Artem Ploujnikov 2021
11+
"""
12+
13+
import torch
14+
from torch import nn
15+
16+
17+
class GuidedAttentionLoss(nn.Module):
18+
"""
19+
A loss implementation that forces attention matrices to be
20+
near-diagonal, imposing progressively larger penalties for paying
21+
attention to regions far away from the diagonal). It is useful
22+
for sequence-to-sequence models in which the sequence of outputs
23+
is expected to corrsespond closely to the sequence of inputs,
24+
such as TTS or G2P
25+
26+
https://arxiv.org/abs/1710.08969
27+
28+
The implementation is inspired by the R9Y9 DeepVoice3 model
29+
https://github.com/r9y9/deepvoice3_pytorch
30+
31+
It should be roughly equivalent to it; however, it has been
32+
fully vectorized.
33+
34+
Arguments
35+
---------
36+
sigma:
37+
the guided attention weight
38+
"""
39+
40+
def __init__(self, sigma=0.2):
41+
super().__init__()
42+
self.sigma = sigma
43+
self.weight_factor = 2 * (sigma ** 2)
44+
45+
def forward(
46+
self,
47+
attention,
48+
input_lengths,
49+
target_lengths,
50+
max_input_len=None,
51+
max_target_len=None,
52+
):
53+
"""
54+
Computes the guided attention loss for a single batch
55+
56+
Arguments
57+
---------
58+
attention: torch.Tensor
59+
A padded attention/alignments matrix
60+
(batch, targets, inputs)
61+
input_lengths: torch.tensor
62+
A (batch, lengths) tensor of input lengths
63+
target_lengths: torch.tensor
64+
A (batch, lengths) tensor of target lengths
65+
max_input_len: int
66+
The maximum input length - optional,
67+
if not computed will be set to the maximum
68+
of target_lengths. Setting it explicitly
69+
might be necessary when using data parallelism
70+
max_target_len: int
71+
The maximum target length - optional,
72+
if not computed will be set to the maximum
73+
of target_lengths. Setting it explicitly
74+
might be necessary when using data parallelism
75+
76+
77+
Returns
78+
-------
79+
loss: torch.Tensor
80+
A single-element tensor with the loss value
81+
"""
82+
soft_mask = self.guided_attentions(
83+
input_lengths, target_lengths, max_input_len, max_target_len
84+
)
85+
return (attention * soft_mask.transpose(-1, -2)).mean()
86+
87+
def guided_attentions(
88+
self,
89+
input_lengths,
90+
target_lengths,
91+
max_input_len=None,
92+
max_target_len=None,
93+
):
94+
"""
95+
Computes guided attention matrices
96+
97+
Arguments
98+
---------
99+
input_lengths: torch.Tensor
100+
A tensor of input lengths
101+
target_lengths: torch.Tensor
102+
A tensor of target lengths
103+
max_input_len: int
104+
The maximum input length - optional,
105+
if not computed will be set to the maximum
106+
of target_lengths. Setting it explicitly
107+
might be necessary when using data parallelism
108+
max_target_len: int
109+
The maximum target length - optional,
110+
if not computed will be set to the maximum
111+
of target_lengths. Setting it explicitly
112+
might be necessary when using data parallelism
113+
114+
Returns
115+
-------
116+
soft_mask: torch.Tensor
117+
The guided attention tensor of shape (batch, max_input_len, max_target_len)
118+
"""
119+
input_lengths_broad = input_lengths.view(-1, 1, 1)
120+
target_lengths_broad = target_lengths.view(-1, 1, 1)
121+
if max_input_len is None:
122+
max_input_len = input_lengths.max()
123+
if max_target_len is None:
124+
max_target_len = target_lengths.max()
125+
input_mesh, target_mesh = torch.meshgrid(
126+
torch.arange(max_input_len).to(input_lengths.device),
127+
torch.arange(max_target_len).to(target_lengths.device),
128+
)
129+
input_mesh, target_mesh = (
130+
input_mesh.unsqueeze(0),
131+
target_mesh.unsqueeze(0),
132+
)
133+
input_lengths_broad = input_lengths.view(-1, 1, 1)
134+
target_lengths_broad = target_lengths.view(-1, 1, 1)
135+
soft_mask = 1.0 - torch.exp(
136+
-(
137+
(
138+
input_mesh / input_lengths_broad
139+
- target_mesh / target_lengths_broad
140+
)
141+
** 2
142+
)
143+
/ self.weight_factor
144+
)
145+
outside = (input_mesh >= input_lengths_broad) | (
146+
target_mesh >= target_lengths_broad
147+
)
148+
soft_mask[outside] = 0.0
149+
return soft_mask

‎tests/unittests/test_losses.py‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,3 +137,90 @@ def test_transducer_loss():
137137
)
138138
out_cost.backward()
139139
assert out_cost.item() == 2.247833251953125
140+
141+
142+
def test_guided_attention_loss_mask():
143+
from speechbrain.nnet.loss.guidedattn_loss import GuidedAttentionLoss
144+
145+
loss = GuidedAttentionLoss()
146+
input_lengths = torch.tensor([3, 2, 6])
147+
output_lengths = torch.tensor([4, 3, 5])
148+
soft_mask = loss.guided_attentions(input_lengths, output_lengths)
149+
ref_soft_mask = torch.tensor(
150+
[
151+
[
152+
[0.0, 0.54216665, 0.9560631, 0.9991162, 0.0],
153+
[0.7506478, 0.08314464, 0.2933517, 0.8858382, 0.0],
154+
[0.9961341, 0.8858382, 0.2933517, 0.08314464, 0.0],
155+
[0.0, 0.0, 0.0, 0.0, 0.0],
156+
[0.0, 0.0, 0.0, 0.0, 0.0],
157+
[0.0, 0.0, 0.0, 0.0, 0.0],
158+
],
159+
[
160+
[0.0, 0.7506478, 0.9961341, 0.0, 0.0],
161+
[0.9560631, 0.2933517, 0.2933517, 0.0, 0.0],
162+
[0.0, 0.0, 0.0, 0.0, 0.0],
163+
[0.0, 0.0, 0.0, 0.0, 0.0],
164+
[0.0, 0.0, 0.0, 0.0, 0.0],
165+
[0.0, 0.0, 0.0, 0.0, 0.0],
166+
],
167+
[
168+
[0.0, 0.39346933, 0.86466473, 0.988891, 0.99966455],
169+
[0.2933517, 0.01379288, 0.49366438, 0.90436554, 0.993355],
170+
[0.7506478, 0.1992626, 0.05404053, 0.5888877, 0.93427145],
171+
[0.9560631, 0.6753475, 0.1175031, 0.1175031, 0.6753475],
172+
[0.9961341, 0.93427145, 0.5888877, 0.05404053, 0.1992626],
173+
[0.9998301, 0.993355, 0.90436554, 0.49366438, 0.01379288],
174+
],
175+
]
176+
)
177+
assert torch.allclose(soft_mask, ref_soft_mask)
178+
179+
180+
def test_guided_attention_loss_value():
181+
from speechbrain.nnet.loss.guidedattn_loss import GuidedAttentionLoss
182+
183+
loss = GuidedAttentionLoss()
184+
input_lengths = torch.tensor([2, 3])
185+
target_lengths = torch.tensor([3, 4])
186+
alignments = torch.tensor(
187+
[
188+
[
189+
[0.8, 0.2, 0.0],
190+
[0.4, 0.6, 0.0],
191+
[0.2, 0.8, 0.0],
192+
[0.0, 0.0, 0.0],
193+
],
194+
[
195+
[0.6, 0.2, 0.2],
196+
[0.1, 0.7, 0.2],
197+
[0.3, 0.4, 0.3],
198+
[0.2, 0.3, 0.5],
199+
],
200+
]
201+
)
202+
loss_value = loss(alignments, input_lengths, target_lengths)
203+
ref_loss_value = torch.tensor(0.1142)
204+
assert torch.isclose(loss_value, ref_loss_value, 0.0001, 0.0001).item()
205+
206+
207+
def test_guided_attention_loss_shapes():
208+
from speechbrain.nnet.loss.guidedattn_loss import GuidedAttentionLoss
209+
210+
loss = GuidedAttentionLoss()
211+
input_lengths = torch.tensor([3, 2, 6])
212+
output_lengths = torch.tensor([4, 3, 5])
213+
soft_mask = loss.guided_attentions(input_lengths, output_lengths)
214+
assert soft_mask.shape == (3, 6, 5)
215+
soft_mask = loss.guided_attentions(
216+
input_lengths, output_lengths, max_input_len=10
217+
)
218+
assert soft_mask.shape == (3, 10, 5)
219+
soft_mask = loss.guided_attentions(
220+
input_lengths, output_lengths, max_target_len=12
221+
)
222+
assert soft_mask.shape == (3, 6, 12)
223+
soft_mask = loss.guided_attentions(
224+
input_lengths, output_lengths, max_input_len=10, max_target_len=12
225+
)
226+
assert soft_mask.shape == (3, 10, 12)

0 commit comments

Comments
 (0)