88 lines
2.5 KiB
Python
88 lines
2.5 KiB
Python
import os
|
|
import sys
|
|
import unittest
|
|
import torch
|
|
|
|
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
|
sys.path.append(MODULE_ROOT)
|
|
from pytorch_wrapper.ptn_gru_linear import PtnGRULinear
|
|
|
|
|
|
class TestPtnGRULinear(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Set up test instance."""
|
|
self.model_node = PtnGRULinear()
|
|
|
|
def test_seq_first_1_layer_uni(self):
|
|
self.model = self.model_node.f(
|
|
2, # input_size
|
|
3, # hidden_size
|
|
1, # num_layers
|
|
True, # bias,
|
|
False, # batch_first so [Seq, B, Token]
|
|
0.0, # dropout
|
|
False, # bidirectional,
|
|
5, # linear output features
|
|
True # linear bias
|
|
)[0]
|
|
|
|
x = torch.randn(10, 8, 2) # [Seq, B, Token]
|
|
outputs= self.model(x)
|
|
self.assertEqual(outputs.size(), torch.Size([8, 5]))
|
|
|
|
def test_seq_first_1_layer_bidi(self):
|
|
self.model = self.model_node.f(
|
|
2, # input_size
|
|
3, # hidden_size
|
|
1, # num_layers
|
|
True, # bias,
|
|
False, # batch_first so [Seq, B, Token]
|
|
0.0, # dropout
|
|
True, # bidirectional,
|
|
5, # linear output features
|
|
True # linear bias
|
|
)[0]
|
|
|
|
x = torch.randn(10, 8, 2) # [Seq, B, Token]
|
|
outputs= self.model(x)
|
|
self.assertEqual(outputs.size(), torch.Size([8, 5]))
|
|
|
|
def test_seq_first_2_layer_uni(self):
|
|
self.model = self.model_node.f(
|
|
2, # input_size
|
|
3, # hidden_size
|
|
2, # num_layers
|
|
True, # bias,
|
|
False, # batch_first so [Seq, B, Token]
|
|
0.0, # dropout
|
|
False, # bidirectional,
|
|
5, # linear output features
|
|
True # linear bias
|
|
)[0]
|
|
|
|
x = torch.randn(10, 8, 2) # [Seq, B, Token]
|
|
outputs= self.model(x)
|
|
self.assertEqual(outputs.size(), torch.Size([8, 5]))
|
|
|
|
def test_batch_first_1_layer_uni(self):
|
|
self.model = self.model_node.f(
|
|
2, # input_size
|
|
3, # hidden_size
|
|
1, # num_layers
|
|
True, # bias,
|
|
True, # batch_first
|
|
0.0, # dropout
|
|
False, # bidirectional,
|
|
5, # linear output features
|
|
True # linear bias
|
|
)[0]
|
|
|
|
x = torch.randn(8, 10, 2) # [B, Seq, Token]
|
|
outputs= self.model(x)
|
|
self.assertEqual(outputs.size(), torch.Size([8, 5]))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|