Add test code for Python implementation of transformer

README.md

@@ -31,6 +31,16 @@ python ./experiments/classify.py  --random-seed=1234 --num-epochs=1 --tiny
 
 ```
 
+# Tests
+
+## Python implementation tests
+To run the tests for the Python implementation of the transformer, execute the following command:
+
+```
+cd src/python
+python -m unittest test_transformer.py
+```
+
 # TODO:
 - main Makefile for build library and c executable
 - add config load

src/c/test_transformer.c

@@ -0,0 +1,64 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <assert.h>
+#include "transformer.h"
+
+void test_self_attention() {
+    struct Transformer *trfm = init_transformer();
+    assert(trfm);
+
+    Tensor x, z;
+    x = MatDim(trfm->n_seq, trfm->dim_embedding);
+    MatFill(x, 4);
+
+    z = self_attention(&trfm->sa[0], x);
+
+    assert(MatRows(z) == trfm->n_seq);
+    assert(MatCols(z) == trfm->dim_internal);
+
+    MatUnDim(x);
+    MatUnDim(z);
+}
+
+void test_multi_head_attention() {
+    struct Transformer *trfm = init_transformer();
+    assert(trfm);
+
+    Tensor x, z;
+    x = MatDim(trfm->n_seq, trfm->dim_embedding);
+    MatFill(x, 4);
+
+    z = multi_head_attention(trfm, x);
+
+    assert(MatRows(z) == trfm->n_seq);
+    assert(MatCols(z) == trfm->dim_embedding);
+
+    MatUnDim(x);
+    MatUnDim(z);
+}
+
+void test_transformer_block() {
+    struct Transformer *trfm = init_transformer();
+    assert(trfm);
+
+    Tensor x, z;
+    x = MatDim(trfm->n_seq, trfm->dim_embedding);
+    MatFill(x, 4);
+
+    z = transformer_block(trfm, x);
+
+    assert(MatRows(z) == trfm->n_seq);
+    assert(MatCols(z) == trfm->dim_embedding);
+
+    MatUnDim(x);
+    MatUnDim(z);
+}
+
+int main() {
+    test_self_attention();
+    test_multi_head_attention();
+    test_transformer_block();
+
+    printf("All tests passed.\n");
+    return 0;
+}

src/python/test_transformer.py

@@ -0,0 +1,50 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import unittest
+from transformer_simple import SelfAttention_naive, MultiHeadAttention_naive, TransformerBlock_naive
+
+class TestSelfAttentionNaive(unittest.TestCase):
+    def setUp(self):
+        self.dim_emb = 4
+        self.dim_internal = 3
+        self.heads = 8
+        self.n_seq = 2
+        self.dtype = torch.float32
+        self.model = SelfAttention_naive(self.dim_emb, self.dim_internal, self.heads, dtype=self.dtype)
+        self.x = torch.ones([1, self.n_seq, self.dim_emb], dtype=self.dtype)
+
+    def test_forward(self):
+        output = self.model(self.x)
+        self.assertEqual(output.shape, (1, self.n_seq, self.dim_internal))
+
+class TestMultiHeadAttentionNaive(unittest.TestCase):
+    def setUp(self):
+        self.dim_emb = 4
+        self.dim_internal = 3
+        self.heads = 8
+        self.n_seq = 2
+        self.dtype = torch.float32
+        self.model = MultiHeadAttention_naive(self.n_seq, self.dim_emb, self.dim_internal, self.heads, dtype=self.dtype)
+        self.x = torch.ones([1, self.n_seq, self.dim_emb], dtype=self.dtype)
+
+    def test_forward(self):
+        output = self.model(self.x)
+        self.assertEqual(output.shape, (1, self.n_seq, self.dim_emb))
+
+class TestTransformerBlockNaive(unittest.TestCase):
+    def setUp(self):
+        self.dim_emb = 4
+        self.dim_internal = 3
+        self.heads = 8
+        self.n_seq = 2
+        self.dtype = torch.float32
+        self.model = TransformerBlock_naive(self.n_seq, self.dim_emb, self.dim_internal, self.heads, dtype=self.dtype)
+        self.x = torch.ones([1, self.n_seq, self.dim_emb], dtype=self.dtype)
+
+    def test_forward(self):
+        output = self.model(self.x)
+        self.assertEqual(output.shape, (1, self.n_seq, self.dim_emb))
+
+if __name__ == '__main__':
+    unittest.main()