mirror of
https://github.com/facebookresearch/fairseq.git
synced 2024-08-16 20:10:40 +03:00
177 lines
5.4 KiB
Python
177 lines
5.4 KiB
Python
|
import torch
|
||
|
import numpy as np
|
||
|
import unittest
|
||
|
from fairseq.modules import (
|
||
|
ESPNETMultiHeadedAttention,
|
||
|
RelPositionMultiHeadedAttention,
|
||
|
RotaryPositionMultiHeadedAttention,
|
||
|
)
|
||
|
|
||
|
torch.use_deterministic_algorithms(True)
|
||
|
|
||
|
|
||
|
class TestESPNETMultiHeadedAttention(unittest.TestCase):
|
||
|
def setUp(self) -> None:
|
||
|
self.T = 3
|
||
|
self.B = 1
|
||
|
self.C = 2
|
||
|
torch.manual_seed(0)
|
||
|
self.sample = torch.randn(self.T, self.B, self.C) # TBC
|
||
|
self.sample_scores = torch.randn(self.B, 1, self.T, self.T)
|
||
|
self.MHA = ESPNETMultiHeadedAttention(self.C, 1, dropout=0)
|
||
|
|
||
|
def test_forward(self):
|
||
|
expected_scores = torch.tensor(
|
||
|
[[[0.1713, -0.3776]], [[0.2263, -0.4486]], [[0.2243, -0.4538]]]
|
||
|
)
|
||
|
scores, _ = self.MHA(self.sample, self.sample, self.sample)
|
||
|
self.assertTrue(
|
||
|
np.allclose(
|
||
|
expected_scores.cpu().detach().numpy(),
|
||
|
scores.cpu().detach().numpy(),
|
||
|
atol=1e-4,
|
||
|
)
|
||
|
)
|
||
|
|
||
|
def test_forward_qkv(self):
|
||
|
expected_query = torch.tensor(
|
||
|
[[[[-1.0235, 0.0409], [0.4008, 1.3077], [0.5396, 2.0698]]]]
|
||
|
)
|
||
|
expected_key = torch.tensor(
|
||
|
[[[[0.5053, -0.4965], [-0.3730, -0.9473], [-0.7019, -0.1935]]]]
|
||
|
)
|
||
|
expected_val = torch.tensor(
|
||
|
[[[[-0.9940, 0.5403], [0.5924, -0.7619], [0.7504, -1.0892]]]]
|
||
|
)
|
||
|
sample_t = self.sample.transpose(0, 1)
|
||
|
query, key, val = self.MHA.forward_qkv(sample_t, sample_t, sample_t)
|
||
|
self.assertTrue(
|
||
|
np.allclose(
|
||
|
expected_query.cpu().detach().numpy(),
|
||
|
query.cpu().detach().numpy(),
|
||
|
atol=1e-4,
|
||
|
)
|
||
|
)
|
||
|
self.assertTrue(
|
||
|
np.allclose(
|
||
|
expected_key.cpu().detach().numpy(),
|
||
|
key.cpu().detach().numpy(),
|
||
|
atol=1e-4,
|
||
|
)
|
||
|
)
|
||
|
self.assertTrue(
|
||
|
np.allclose(
|
||
|
expected_val.cpu().detach().numpy(),
|
||
|
val.cpu().detach().numpy(),
|
||
|
atol=1e-4,
|
||
|
)
|
||
|
)
|
||
|
|
||
|
def test_forward_attention(self):
|
||
|
expected_scores = torch.tensor(
|
||
|
[[[0.1627, -0.6249], [-0.2547, -0.6487], [-0.0711, -0.8545]]]
|
||
|
)
|
||
|
scores = self.MHA.forward_attention(
|
||
|
self.sample.transpose(0, 1).view(self.B, 1, self.T, self.C),
|
||
|
self.sample_scores,
|
||
|
mask=None,
|
||
|
)
|
||
|
self.assertTrue(
|
||
|
np.allclose(
|
||
|
expected_scores.cpu().detach().numpy(),
|
||
|
scores.cpu().detach().numpy(),
|
||
|
atol=1e-4,
|
||
|
)
|
||
|
)
|
||
|
|
||
|
|
||
|
class TestRelPositionMultiHeadedAttention(unittest.TestCase):
|
||
|
def setUp(self) -> None:
|
||
|
self.T = 3
|
||
|
self.B = 1
|
||
|
self.C = 2
|
||
|
torch.manual_seed(0)
|
||
|
self.sample = torch.randn(self.T, self.B, self.C) # TBC
|
||
|
self.sample_x = torch.randn(self.B, 1, self.T, self.T * 2 - 1)
|
||
|
self.sample_pos = torch.randn(self.B, self.T * 2 - 1, self.C)
|
||
|
self.MHA = RelPositionMultiHeadedAttention(self.C, 1, dropout=0)
|
||
|
|
||
|
def test_rel_shift(self):
|
||
|
expected_x = torch.tensor(
|
||
|
[
|
||
|
[
|
||
|
[
|
||
|
[-0.7193, -0.4033, -0.5966],
|
||
|
[-0.8567, 1.1006, -1.0712],
|
||
|
[-0.5663, 0.3731, -0.8920],
|
||
|
]
|
||
|
]
|
||
|
]
|
||
|
)
|
||
|
x = self.MHA.rel_shift(self.sample_x)
|
||
|
self.assertTrue(
|
||
|
np.allclose(
|
||
|
expected_x.cpu().detach().numpy(),
|
||
|
x.cpu().detach().numpy(),
|
||
|
atol=1e-4,
|
||
|
)
|
||
|
)
|
||
|
|
||
|
def test_forward(self):
|
||
|
expected_scores = torch.tensor(
|
||
|
[
|
||
|
[[-0.9609, -0.5020]],
|
||
|
[[-0.9308, -0.4890]],
|
||
|
[[-0.9473, -0.4948]],
|
||
|
[[-0.9609, -0.5020]],
|
||
|
[[-0.9308, -0.4890]],
|
||
|
[[-0.9473, -0.4948]],
|
||
|
[[-0.9609, -0.5020]],
|
||
|
[[-0.9308, -0.4890]],
|
||
|
[[-0.9473, -0.4948]],
|
||
|
[[-0.9609, -0.5020]],
|
||
|
[[-0.9308, -0.4890]],
|
||
|
[[-0.9473, -0.4948]],
|
||
|
[[-0.9609, -0.5020]],
|
||
|
[[-0.9308, -0.4890]],
|
||
|
[[-0.9473, -0.4948]],
|
||
|
]
|
||
|
)
|
||
|
scores, _ = self.MHA(self.sample, self.sample, self.sample, self.sample_pos)
|
||
|
self.assertTrue(
|
||
|
np.allclose(
|
||
|
expected_scores.cpu().detach().numpy(),
|
||
|
scores.cpu().detach().numpy(),
|
||
|
atol=1e-4,
|
||
|
)
|
||
|
)
|
||
|
|
||
|
|
||
|
class TestRotaryPositionMultiHeadedAttention(unittest.TestCase):
|
||
|
def setUp(self) -> None:
|
||
|
self.T = 3
|
||
|
self.B = 1
|
||
|
self.C = 2
|
||
|
torch.manual_seed(0)
|
||
|
self.sample = torch.randn(self.T, self.B, self.C) # TBC
|
||
|
self.MHA = RotaryPositionMultiHeadedAttention(
|
||
|
self.C, 1, dropout=0, precision=None
|
||
|
)
|
||
|
|
||
|
def test_forward(self):
|
||
|
expected_scores = torch.tensor(
|
||
|
[[[-0.3220, -0.4726]], [[-1.2813, -0.0979]], [[-0.3138, -0.4758]]]
|
||
|
)
|
||
|
scores, _ = self.MHA(self.sample, self.sample, self.sample)
|
||
|
self.assertTrue(
|
||
|
np.allclose(
|
||
|
expected_scores.cpu().detach().numpy(),
|
||
|
scores.cpu().detach().numpy(),
|
||
|
atol=1e-4,
|
||
|
)
|
||
|
)
|
||
|
|
||
|
|
||
|
if __name__ == "__main__":
|
||
|
unittest.main()
|