add star-transformer
@@ -3,6 +3,7 @@ import unittest
import torch
from fastNLP.modules.other_modules import GroupNorm, LayerNormalization, BiLinear, BiAffine
+from fastNLP.modules.encoder.star_transformer import StarTransformer
class TestGroupNorm(unittest.TestCase):
@@ -49,3 +50,12 @@ class TestBiAffine(unittest.TestCase):
encoder_input = torch.randn((batch_size, decoder_length, 10))
y = layer(decoder_input, encoder_input)
self.assertEqual(tuple(y.shape), (batch_size, 25, encoder_length, 1))
+
+class TestStarTransformer(unittest.TestCase):
+ def test_1(self):
+ model = StarTransformer(num_layers=6, hidden_size=100, num_head=8, head_dim=20, max_len=100)
+ x = torch.rand(16, 45, 100)
+ mask = torch.ones(16, 45).byte()
+ y, yn = model(x, mask)
+ self.assertEqual(tuple(y.size()), (16, 45, 100))
+ self.assertEqual(tuple(yn.size()), (16, 100))
