<filename>pytorch_translate/test/test_attention.py
#!/usr/bin/env python3

import random
import unittest

import numpy as np
import pytorch_translate.attention.multihead_attention as multihead_attention
import torch
from pytorch_translate.attention import attention_utils, dot_attention, mlp_attention


class TestAttention(unittest.TestCase):
    def setUp(self):
        self.bsz = 10
        self.src_len = 5
        self.ctx_dim = 3
        self.dec_dim = 4
        self.att_dim = 2

    def test_masked_softmax(self):
        scores = torch.rand(20, 20)
        lengths = torch.arange(start=1, end=21)

        masked_normalized_scores = attention_utils.masked_softmax(
            scores, lengths, src_length_masking=True
        )

        for i in range(20):
            scores_sum = masked_normalized_scores[i].numpy().sum()
            self.assertAlmostEqual(scores_sum, 1, places=6)

    def _test_attention(self, attention):
        dummy_source_hids = torch.rand(self.src_len, self.bsz, self.ctx_dim)
        dummy_decoder_state = torch.rand(self.bsz, self.dec_dim)
        dummy_src_lengths = torch.fmod(torch.arange(self.bsz), self.src_len) + 1
        attention(dummy_decoder_state, dummy_source_hids, dummy_src_lengths)

    def test_dot_attention(self):
        self._test_attention(
            dot_attention.DotAttention(
                self.dec_dim,
                self.ctx_dim,
                src_length_masking=True,
                force_projection=True,
            )
        )

    def test_mlp_attention(self):
        self._test_attention(
            mlp_attention.MLPAttention(
                self.dec_dim,
                self.ctx_dim,
                src_length_masking=True,
                attention_dim=self.att_dim,
            )
        )


def _softmax(x):  # softmax over 4 dim matrix
    """ Numpy-based reference softmax over 4 dim matrix"""
    output = np.zeros(x.shape, dtype=np.float32)
    for i in range(x.shape[0]):
        for j in range(x.shape[1]):
            for k in range(x.shape[2]):
                x_curr = x[i, j, k, :]
                e_x = np.exp(x_curr - np.amax(x_curr))
                output[i, j, k, :] = e_x / np.sum(e_x)
    return output


def _batchmatmul(a, b):  # batchmatmul over 4 dim matrix
    """ Numpy-based batch matrix multiply over 4 dim matrix"""
    assert a.shape[0] == b.shape[0]
    assert a.shape[1] == b.shape[1]
    retval = np.zeros(
        (a.shape[0], a.shape[1], a.shape[2], b.shape[3]), dtype=np.float32
    )
    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
            retval[i, j, :, :] = np.matmul(a[i, j, :, :], b[i, j, :, :])
    return retval


class MultiheadAttentionTest(unittest.TestCase):
    def _scaled_dot_attn_ref(self, Q, K, V, dims, unseen_mask=False, src_lengths=None):
        """ Numpy-based reference implementation of scaled dot attention
        for testing"""
        QKT = _batchmatmul(
            Q,
            np.transpose(K, axes=[0, 1, 3, 2])
            / np.sqrt(dims[3], dtype=np.float32),  # divide by sqrt(d_head)
        )
        if unseen_mask or src_lengths is not None:
            b1, b2, s1, s2 = QKT.shape
            # assert s1 == s2
            for i in range(b1):
                for j in range(b2):
                    for m in range(s1):
                        for n in range(s2):
                            if unseen_mask and n > m:
                                QKT[i, j, m, n] = -np.inf
                            if src_lengths is not None and n >= src_lengths[i]:
                                QKT[i, j, m, n] = -np.inf
        reference = _softmax(QKT)
        reference = _batchmatmul(reference, V)
        return reference

    def _generate_src_lengths(self, batch_size, seq_len):
        src_lengths = np.array([random.randint(1, seq_len) for i in range(batch_size)])

        # max source length has to equal seq_len, so randomly choose
        # one example to have source length = seq_len
        max_len_example_i = random.randint(0, batch_size - 1)
        src_lengths[max_len_example_i] = seq_len

        src_lengths_tensor = torch.from_numpy(src_lengths).int()
        return src_lengths, src_lengths_tensor

    def _split_heads_ref(self, X, dims, nheads, d_head):
        X_split = np.reshape(X, dims[:2] + [nheads, d_head])
        X_split_transposed = np.transpose(X_split, [0, 2, 1, 3])
        reference = np.reshape(X_split_transposed, [dims[0], nheads, dims[1], d_head])
        return reference

    def _combine_heads_ref(self, X, dims, nheads, d_head):
        X_transposed = np.transpose(X, [0, 2, 1, 3])
        reference = np.reshape(X_transposed, dims[:2] + [nheads * d_head])
        return reference

    def _fc(self, X, X_name, module, start=None, end=None):
        X_fc_b = None
        X_fc_w = None
        for name, param in module.named_parameters():
            if X_name + ".weight" in name or X_name + "_weight" in name:
                if X_fc_w is not None:
                    raise Exception(f"Duplicate FC name {name} found")
                X_fc_w = param.detach().numpy()
            elif X_name + ".bias" in name or X_name + "_bias" in name:
                if X_fc_b is not None:
                    raise Exception(f"Duplicate FC name {name} found")
                X_fc_b = param.detach().numpy()
        return np.matmul(X, np.transpose(X_fc_w)) + X_fc_b

    def _multihead_attn_test_helper(self, use_src_lengths):
        for _ in range(100):
            batch_sz, seq_len = [random.randint(2, 10) for r in range(2)]
            d_head = random.randint(3, 10)
            nheads = random.randint(3, 10)
            d_model = d_head * nheads
            dims = [batch_sz, seq_len, d_model]

            src_lengths = None
            src_lengths_tensor = None
            if use_src_lengths:
                src_lengths, src_lengths_tensor = self._generate_src_lengths(
                    batch_size=batch_sz, seq_len=seq_len
                )

            decoder_state = np.random.rand(batch_sz, d_model).astype(np.float32)
            K = np.random.rand(*dims).astype(np.float32)
            V = K
            Q = np.expand_dims(decoder_state, 1)

            decoder_state_tensor = torch.from_numpy(decoder_state).float()
            source_hid_tensor = torch.from_numpy(K).float().transpose(0, 1)

            multihead_attn_module = multihead_attention.MultiheadAttention(
                context_dim=d_model, decoder_hidden_state_dim=d_model, nheads=nheads
            )

            for squeeze in [True, False]:
                result = (
                    multihead_attn_module(
                        decoder_state=decoder_state_tensor,
                        source_hids=source_hid_tensor,
                        src_lengths=src_lengths_tensor,
                        squeeze=squeeze,
                    )[0]
                    .detach()
                    .numpy()
                )

                if not squeeze:
                    self.assertEqual(result.ndim, 3)
                    result = np.squeeze(result, axis=0)

                Q_fc = self._fc(Q, "q_proj", multihead_attn_module, end=d_model)
                K_fc = self._fc(
                    K, "k_proj", multihead_attn_module, start=d_model, end=2 * d_model
                )
                V_fc = self._fc(V, "v_proj", multihead_attn_module, start=2 * d_model)

                Q_split = self._split_heads_ref(
                    Q_fc, [batch_sz, 1, d_model], nheads, d_head
                )
                K_split = self._split_heads_ref(K_fc, dims, nheads, d_head)
                V_split = self._split_heads_ref(V_fc, dims, nheads, d_head)

                attn_heads = self._scaled_dot_attn_ref(
                    Q=Q_split,
                    K=K_split,
                    V=V_split,
                    dims=Q_split.shape,
                    src_lengths=src_lengths,
                )

                combined_attn_heads = self._combine_heads_ref(
                    X=attn_heads, dims=[batch_sz, 1], nheads=nheads, d_head=d_head
                )

                reference = self._fc(
                    combined_attn_heads, "out_proj", multihead_attn_module
                )
                reference = np.squeeze(reference, axis=1)
                self.assertEqual(tuple(result.shape), (batch_sz, d_model))
                np.testing.assert_allclose(result, reference, atol=1e-5)

    def test_multihead_attn_no_masking(self):
        self._multihead_attn_test_helper(use_src_lengths=None)

    def test_multihead_attn_with_src_lengths(self):
        self._multihead_attn_test_helper(use_src_lengths=True)

<filename>pytorch_translate/test/test_attention.py
#!/usr/bin/env python3

import random
import unittest

import numpy as np
import pytorch_translate.attention.multihead_attention as multihead_attention
import torch
from pytorch_translate.attention import attention_utils, dot_attention, mlp_attention


class TestAttention(unittest.TestCase):
    def setUp(self):
        self.bsz = 10
        self.src_len = 5
        self.ctx_dim = 3
        self.dec_dim = 4
        self.att_dim = 2

    def test_masked_softmax(self):
        scores = torch.rand(20, 20)
        lengths = torch.arange(start=1, end=21)

        masked_normalized_scores = attention_utils.masked_softmax(
            scores, lengths, src_length_masking=True
        )

        for i in range(20):
            scores_sum = masked_normalized_scores[i].numpy().sum()
            self.assertAlmostEqual(scores_sum, 1, places=6)

    def _test_attention(self, attention):
        dummy_source_hids = torch.rand(self.src_len, self.bsz, self.ctx_dim)
        dummy_decoder_state = torch.rand(self.bsz, self.dec_dim)
        dummy_src_lengths = torch.fmod(torch.arange(self.bsz), self.src_len) + 1
        attention(dummy_decoder_state, dummy_source_hids, dummy_src_lengths)

    def test_dot_attention(self):
        self._test_attention(
            dot_attention.DotAttention(
                self.dec_dim,
                self.ctx_dim,
                src_length_masking=True,
                force_projection=True,
            )
        )

    def test_mlp_attention(self):
        self._test_attention(
            mlp_attention.MLPAttention(
                self.dec_dim,
                self.ctx_dim,
                src_length_masking=True,
                attention_dim=self.att_dim,
            )
        )


def _softmax(x):  # softmax over 4 dim matrix
    """ Numpy-based reference softmax over 4 dim matrix"""
    output = np.zeros(x.shape, dtype=np.float32)
    for i in range(x.shape[0]):
        for j in range(x.shape[1]):
            for k in range(x.shape[2]):
                x_curr = x[i, j, k, :]
                e_x = np.exp(x_curr - np.amax(x_curr))
                output[i, j, k, :] = e_x / np.sum(e_x)
    return output


def _batchmatmul(a, b):  # batchmatmul over 4 dim matrix
    """ Numpy-based batch matrix multiply over 4 dim matrix"""
    assert a.shape[0] == b.shape[0]
    assert a.shape[1] == b.shape[1]
    retval = np.zeros(
        (a.shape[0], a.shape[1], a.shape[2], b.shape[3]), dtype=np.float32
    )
    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
            retval[i, j, :, :] = np.matmul(a[i, j, :, :], b[i, j, :, :])
    return retval


class MultiheadAttentionTest(unittest.TestCase):
    def _scaled_dot_attn_ref(self, Q, K, V, dims, unseen_mask=False, src_lengths=None):
        """ Numpy-based reference implementation of scaled dot attention
        for testing"""
        QKT = _batchmatmul(
            Q,
            np.transpose(K, axes=[0, 1, 3, 2])
            / np.sqrt(dims[3], dtype=np.float32),  # divide by sqrt(d_head)
        )
        if unseen_mask or src_lengths is not None:
            b1, b2, s1, s2 = QKT.shape
            # assert s1 == s2
            for i in range(b1):
                for j in range(b2):
                    for m in range(s1):
                        for n in range(s2):
                            if unseen_mask and n > m:
                                QKT[i, j, m, n] = -np.inf
                            if src_lengths is not None and n >= src_lengths[i]:
                                QKT[i, j, m, n] = -np.inf
        reference = _softmax(QKT)
        reference = _batchmatmul(reference, V)
        return reference

    def _generate_src_lengths(self, batch_size, seq_len):
        src_lengths = np.array([random.randint(1, seq_len) for i in range(batch_size)])

        # max source length has to equal seq_len, so randomly choose
        # one example to have source length = seq_len
        max_len_example_i = random.randint(0, batch_size - 1)
        src_lengths[max_len_example_i] = seq_len

        src_lengths_tensor = torch.from_numpy(src_lengths).int()
        return src_lengths, src_lengths_tensor

    def _split_heads_ref(self, X, dims, nheads, d_head):
        X_split = np.reshape(X, dims[:2] + [nheads, d_head])
        X_split_transposed = np.transpose(X_split, [0, 2, 1, 3])
        reference = np.reshape(X_split_transposed, [dims[0], nheads, dims[1], d_head])
        return reference

    def _combine_heads_ref(self, X, dims, nheads, d_head):
        X_transposed = np.transpose(X, [0, 2, 1, 3])
        reference = np.reshape(X_transposed, dims[:2] + [nheads * d_head])
        return reference

    def _fc(self, X, X_name, module, start=None, end=None):
        X_fc_b = None
        X_fc_w = None
        for name, param in module.named_parameters():
            if X_name + ".weight" in name or X_name + "_weight" in name:
                if X_fc_w is not None:
                    raise Exception(f"Duplicate FC name {name} found")
                X_fc_w = param.detach().numpy()
            elif X_name + ".bias" in name or X_name + "_bias" in name:
                if X_fc_b is not None:
                    raise Exception(f"Duplicate FC name {name} found")
                X_fc_b = param.detach().numpy()
        return np.matmul(X, np.transpose(X_fc_w)) + X_fc_b

    def _multihead_attn_test_helper(self, use_src_lengths):
        for _ in range(100):
            batch_sz, seq_len = [random.randint(2, 10) for r in range(2)]
            d_head = random.randint(3, 10)
            nheads = random.randint(3, 10)
            d_model = d_head * nheads
            dims = [batch_sz, seq_len, d_model]

            src_lengths = None
            src_lengths_tensor = None
            if use_src_lengths:
                src_lengths, src_lengths_tensor = self._generate_src_lengths(
                    batch_size=batch_sz, seq_len=seq_len
                )

            decoder_state = np.random.rand(batch_sz, d_model).astype(np.float32)
            K = np.random.rand(*dims).astype(np.float32)
            V = K
            Q = np.expand_dims(decoder_state, 1)

            decoder_state_tensor = torch.from_numpy(decoder_state).float()
            source_hid_tensor = torch.from_numpy(K).float().transpose(0, 1)

            multihead_attn_module = multihead_attention.MultiheadAttention(
                context_dim=d_model, decoder_hidden_state_dim=d_model, nheads=nheads
            )

            for squeeze in [True, False]:
                result = (
                    multihead_attn_module(
                        decoder_state=decoder_state_tensor,
                        source_hids=source_hid_tensor,
                        src_lengths=src_lengths_tensor,
                        squeeze=squeeze,
                    )[0]
                    .detach()
                    .numpy()
                )

                if not squeeze:
                    self.assertEqual(result.ndim, 3)
                    result = np.squeeze(result, axis=0)

                Q_fc = self._fc(Q, "q_proj", multihead_attn_module, end=d_model)
                K_fc = self._fc(
                    K, "k_proj", multihead_attn_module, start=d_model, end=2 * d_model
                )
                V_fc = self._fc(V, "v_proj", multihead_attn_module, start=2 * d_model)

                Q_split = self._split_heads_ref(
                    Q_fc, [batch_sz, 1, d_model], nheads, d_head
                )
                K_split = self._split_heads_ref(K_fc, dims, nheads, d_head)
                V_split = self._split_heads_ref(V_fc, dims, nheads, d_head)

                attn_heads = self._scaled_dot_attn_ref(
                    Q=Q_split,
                    K=K_split,
                    V=V_split,
                    dims=Q_split.shape,
                    src_lengths=src_lengths,
                )

                combined_attn_heads = self._combine_heads_ref(
                    X=attn_heads, dims=[batch_sz, 1], nheads=nheads, d_head=d_head
                )

                reference = self._fc(
                    combined_attn_heads, "out_proj", multihead_attn_module
                )
                reference = np.squeeze(reference, axis=1)
                self.assertEqual(tuple(result.shape), (batch_sz, d_model))
                np.testing.assert_allclose(result, reference, atol=1e-5)

    def test_multihead_attn_no_masking(self):
        self._multihead_attn_test_helper(use_src_lengths=None)

    def test_multihead_attn_with_src_lengths(self):
        self._multihead_attn_test_helper(use_src_lengths=True)

#!/usr/bin/env python3 # softmax over 4 dim matrix Numpy-based reference softmax over 4 dim matrix # batchmatmul over 4 dim matrix Numpy-based batch matrix multiply over 4 dim matrix Numpy-based reference implementation of scaled dot attention         for testing # divide by sqrt(d_head) # assert s1 == s2 # max source length has to equal seq_len, so randomly choose # one example to have source length = seq_len