Skip to content

Commit 00ca165

Browse files
stashuk-olekfacebook-github-bot
stashuk-olek
authored and
facebook-github-bot
committed
Remove attentions from TransformerOutput and stop computing attention weights in FLAVA (facebookresearch#535)
Summary: The `attentions` field on `TransformerOutput` and `return_attn_weights`/`head_mask` parameters in the FLAVA encoder stack were never used by any consumer. This diffs cleans it up. Later the intent is to simplify attention usage / use common API for them. Reviewed By: OmarPavel Differential Revision: D92927086
1 parent 3c2a85a commit 00ca165

13 files changed

Lines changed: 35 additions & 194 deletions

File tree

‎tests/models/flava/test_flava.py‎

Lines changed: 6 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -216,9 +216,7 @@ def test_forward_image_text(self, image_encoder, text_encoder, flava, inputs):
216216
image, _, text, _ = inputs
217217
actual = flava(image, text)
218218
expected_image = image_encoder(image)
219-
expected_text = text_encoder(
220-
text, return_attn_weights=True, return_hidden_states=True
221-
)
219+
expected_text = text_encoder(text, return_hidden_states=True)
222220
assert actual.text_masked == TransformerOutput()
223221
assert actual.multimodal_masked == TransformerOutput()
224222
assert actual.multimodal == TransformerOutput()
@@ -244,12 +242,8 @@ def test_forward_masked_image_and_text(
244242
)
245243
expected_image = image_encoder(image)
246244
expected_image_masked = image_encoder(image, masked_image)
247-
expected_text = text_encoder(
248-
text, return_attn_weights=True, return_hidden_states=True
249-
)
250-
expected_text_masked = text_encoder(
251-
masked_text, return_attn_weights=True, return_hidden_states=True
252-
)
245+
expected_text = text_encoder(text, return_hidden_states=True)
246+
expected_text_masked = text_encoder(masked_text, return_hidden_states=True)
253247
assert actual.multimodal == TransformerOutput()
254248
assert_expected(actual.text_masked, expected_text_masked)
255249
assert_expected(
@@ -277,9 +271,7 @@ def test_forward_masked_text(self, text_encoder, flava, inputs):
277271
text = torch.ones(2, 3, dtype=torch.int32)
278272
masked_text = torch.ones(2, 3, dtype=torch.int32)
279273
actual = flava(text=text, text_masked=masked_text)
280-
expected_text = text_encoder(
281-
text, return_attn_weights=True, return_hidden_states=True
282-
)
274+
expected_text = text_encoder(text, return_hidden_states=True)
283275

284276
assert actual.multimodal_masked == TransformerOutput()
285277
assert actual.multimodal == TransformerOutput()
@@ -289,9 +281,7 @@ def test_forward_masked_text(self, text_encoder, flava, inputs):
289281
assert_expected(actual.text, expected_text)
290282
assert_expected(
291283
actual.text_masked,
292-
text_encoder(
293-
masked_text, return_attn_weights=True, return_hidden_states=True
294-
),
284+
text_encoder(masked_text, return_hidden_states=True),
295285
)
296286
assert_expected(
297287
actual.projected_text_embeddings, expected_text.last_hidden_state[:, 0, :]
@@ -300,9 +290,7 @@ def test_forward_masked_text(self, text_encoder, flava, inputs):
300290
def test_forward_text(self, text_encoder, flava, inputs):
301291
_, _, text, _ = inputs
302292
actual = flava(text=text)
303-
expected_text = text_encoder(
304-
text, return_attn_weights=True, return_hidden_states=True
305-
)
293+
expected_text = text_encoder(text, return_hidden_states=True)
306294
assert actual.multimodal_masked == TransformerOutput()
307295
assert actual.multimodal == TransformerOutput()
308296
assert actual.image == TransformerOutput()

‎tests/models/flava/test_image_encoder.py‎

Lines changed: 0 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -139,32 +139,3 @@ def test_image_encoder(self, image_encoder_components, input):
139139
atol=1e-4,
140140
rtol=0,
141141
)
142-
assert_expected(
143-
out.attentions,
144-
(
145-
torch.Tensor(
146-
[
147-
[
148-
[
149-
[0.2000, 0.2000, 0.2000, 0.2000, 0.2000],
150-
[0.1999, 0.2000, 0.2000, 0.2000, 0.2000],
151-
[0.1999, 0.2000, 0.2000, 0.2000, 0.2000],
152-
[0.1999, 0.2000, 0.2000, 0.2000, 0.2000],
153-
[0.1999, 0.2000, 0.2000, 0.2000, 0.2000],
154-
]
155-
],
156-
[
157-
[
158-
[0.2000, 0.2000, 0.2000, 0.2000, 0.2000],
159-
[0.1999, 0.2000, 0.2000, 0.2000, 0.2000],
160-
[0.1999, 0.2000, 0.2000, 0.2000, 0.2000],
161-
[0.1999, 0.2000, 0.2000, 0.2000, 0.2000],
162-
[0.1999, 0.2000, 0.2000, 0.2000, 0.2000],
163-
]
164-
],
165-
]
166-
),
167-
),
168-
atol=1e-4,
169-
rtol=0,
170-
)

‎tests/models/flava/test_text_encoder.py‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,6 @@ def test_text_transformer(self, text_encoder_components, input_ids):
7777
text_encoder, _ = text_encoder_components
7878
out = text_encoder(
7979
input_ids,
80-
return_attn_weights=True,
8180
return_hidden_states=True,
8281
)
8382

@@ -95,16 +94,13 @@ def test_text_transformer(self, text_encoder_components, input_ids):
9594
rtol=0.0,
9695
)
9796

98-
assert_expected(out.attentions, (torch.Tensor([[[[0, 1.0], [0.0, 1.0]]]]),))
99-
10097
def test_text_transformer_attn_mask(
10198
self, text_encoder_components, input_ids, attn_mask
10299
):
103100
text_encoder, _ = text_encoder_components
104101
out = text_encoder(
105102
input_ids,
106103
attention_mask=attn_mask,
107-
return_attn_weights=True,
108104
return_hidden_states=True,
109105
)
110106

@@ -123,4 +119,3 @@ def test_text_transformer_attn_mask(
123119
)
124120

125121
assert_expected(out.pooler_output, torch.Tensor([[[1.0, -1.0], [-1.0, 1.0]]]))
126-
assert_expected(out.attentions, (torch.Tensor([[[[1.0, 0], [1.0, 0]]]]),))

‎tests/models/flava/test_transformer.py‎

Lines changed: 2 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -109,11 +109,10 @@ def inputs_ln(self):
109109
return torch.rand((2, 3, 4))
110110

111111
def test_forward(self, inputs, encoder):
112-
output = encoder(inputs, return_hidden_states=True, return_attn_weights=True)
112+
output = encoder(inputs, return_hidden_states=True)
113113

114114
actual_last_hidden_state = output.last_hidden_state
115115
actual_hidden_states = torch.sum(torch.stack(output.hidden_states), dim=0)
116-
actual_attentions = torch.sum(torch.stack(output.attentions), dim=0)
117116

118117
expected_last_hidden_state = torch.Tensor(
119118
[
@@ -127,42 +126,13 @@ def test_forward(self, inputs, encoder):
127126
[[5.1976, 1.9218], [3.8499, 2.2402], [3.1757, -0.1730]],
128127
]
129128
)
130-
expected_attentions = torch.Tensor(
131-
[
132-
[
133-
[
134-
[0.8520, 0.5740, 0.5740],
135-
[0.6232, 0.6884, 0.6884],
136-
[0.6232, 0.6884, 0.6884],
137-
],
138-
[
139-
[0.5859, 0.7071, 0.7071],
140-
[0.6515, 0.6742, 0.6742],
141-
[0.6515, 0.6742, 0.6742],
142-
],
143-
],
144-
[
145-
[
146-
[0.7392, 0.5216, 0.7392],
147-
[0.6434, 0.7132, 0.6434],
148-
[0.7392, 0.5216, 0.7392],
149-
],
150-
[
151-
[0.6207, 0.7586, 0.6207],
152-
[0.6589, 0.6822, 0.6589],
153-
[0.6207, 0.7586, 0.6207],
154-
],
155-
],
156-
]
157-
)
158129

159130
assert_expected(
160131
actual_last_hidden_state, expected_last_hidden_state, rtol=0.0, atol=1e-4
161132
)
162133
assert_expected(
163134
actual_hidden_states, expected_hidden_states, rtol=0.0, atol=1e-4
164135
)
165-
assert_expected(actual_attentions, expected_attentions, rtol=0.0, atol=1e-4)
166136

167137
# set flags to false
168138
output = encoder(inputs)
@@ -172,13 +142,10 @@ def test_forward(self, inputs, encoder):
172142
)
173143

174144
def test_forward_ln(self, inputs_ln, encoder_ln):
175-
output = encoder_ln(
176-
inputs_ln, return_hidden_states=True, return_attn_weights=True
177-
)
145+
output = encoder_ln(inputs_ln, return_hidden_states=True)
178146

179147
actual_last_hidden_state = output.last_hidden_state
180148
actual_hidden_states = torch.sum(torch.stack(output.hidden_states), dim=0)
181-
actual_attentions = torch.sum(torch.stack(output.attentions), dim=0)
182149

183150
expected_last_hidden_state = torch.Tensor(
184151
[
@@ -208,42 +175,13 @@ def test_forward_ln(self, inputs_ln, encoder_ln):
208175
],
209176
]
210177
)
211-
expected_attentions = torch.Tensor(
212-
[
213-
[
214-
[
215-
[0.6653, 0.6376, 0.6971],
216-
[0.7078, 0.5621, 0.7302],
217-
[0.6506, 0.6943, 0.6551],
218-
],
219-
[
220-
[0.6333, 0.7897, 0.5770],
221-
[0.7207, 0.7019, 0.5774],
222-
[0.7285, 0.7195, 0.5520],
223-
],
224-
],
225-
[
226-
[
227-
[0.6919, 0.7021, 0.6060],
228-
[0.6274, 0.7462, 0.6264],
229-
[0.7025, 0.7090, 0.5885],
230-
],
231-
[
232-
[0.5826, 0.6227, 0.7947],
233-
[0.6855, 0.6174, 0.6971],
234-
[0.7317, 0.6057, 0.6625],
235-
],
236-
],
237-
]
238-
)
239178

240179
assert_expected(
241180
actual_last_hidden_state, expected_last_hidden_state, rtol=0.0, atol=1e-4
242181
)
243182
assert_expected(
244183
actual_hidden_states, expected_hidden_states, rtol=0.0, atol=1e-4
245184
)
246-
assert_expected(actual_attentions, expected_attentions, rtol=0.0, atol=1e-4)
247185

248186
# set flags to false
249187
output = encoder_ln(inputs_ln)

‎tests/modules/layers/test_transformer.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,6 @@ def test_forward(
139139
):
140140
assert_expected(state_1, state_2)
141141

142-
assert actual.attentions == expected_output.attentions
143142
assert_expected(
144143
actual.last_hidden_state,
145144
expected_output.last_hidden_state,

‎torchmultimodal/models/albef/multimodal_encoder.py‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -82,9 +82,7 @@ def __init__(
8282
def _self_attention_block(
8383
self, hidden_states: Tensor, attention_mask: Optional[Tensor] = None
8484
) -> Tensor:
85-
output = self.attention(
86-
hidden_states, attention_mask=attention_mask, return_attn_weights=False
87-
)
85+
output = self.attention(hidden_states, attention_mask=attention_mask)
8886
output = self.attention_dropout(output)
8987
return output
9088

@@ -98,7 +96,6 @@ def _cross_attention_block(
9896
hidden_states,
9997
encoder_hidden_states,
10098
attention_mask=cross_attention_mask,
101-
return_attn_weights=False,
10299
)
103100
output = self.cross_attention_dropout(output)
104101
return output

‎torchmultimodal/models/flava/image_encoder.py‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -217,7 +217,6 @@ def forward(
217217
encoder_output = self.encoder(
218218
embedding_output,
219219
attention_mask=attention_mask,
220-
return_attn_weights=True,
221220
return_hidden_states=True,
222221
)
223222
sequence_output = encoder_output.last_hidden_state
@@ -230,7 +229,6 @@ def forward(
230229
last_hidden_state=sequence_output,
231230
pooler_output=pooled_output,
232231
hidden_states=encoder_output.hidden_states,
233-
attentions=encoder_output.attentions,
234232
)
235233

236234

@@ -308,5 +306,4 @@ def forward(
308306
last_hidden_state=output.last_hidden_state,
309307
pooler_output=output.pooler_output,
310308
hidden_states=output.hidden_states,
311-
attentions=output.attentions,
312309
)

‎torchmultimodal/models/flava/model.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,10 @@
77
# Code for some of the transformers components in this file are initialized
88
# from their counterparts in Hugging Face Transformers library.
99

10-
import math
1110
from collections import namedtuple, OrderedDict
1211
from dataclasses import dataclass
1312
from functools import partial
13+
import math
1414
from typing import Any, Callable, List, Optional, Tuple, Union
1515

1616
import torch
@@ -254,7 +254,6 @@ def encode_text(
254254
encoded_text = self.text_encoder(
255255
input_ids=text,
256256
attention_mask=text_mask,
257-
return_attn_weights=True,
258257
return_hidden_states=True,
259258
)
260259
if projection:

0 commit comments

Comments
 (0)