Pytestの使い方

Pytestを使ってみたかったので，使ってみた

Pytestの使い方

今まで研究で実験を行っている中で，テストコードを書いたことがなく，インターンを通じてテストコードの重要を思い知らされたので，これを機に使えるようになっておきたい．

書き方

適当に四則演算をする関数をoperation.pyに用意する．

def add(a, b):
    return a + b

def sub(a, b):
    return a - b

def mul(a, b):
    return a * b

def div(a, b):
    assert b != 0, "b must not be 0!!"
    return a / b

これに対するテストコードをtest_operation.pyに以下のように書く．

from operations import add, div, mul, sub

def test_operations():
    assert add(1, 2) == 3
    assert sub(1, 2) == -1
    assert mul(1, 2) == 2
    assert div(1, 2) == 0.5

pytestでは，test_で始まるファイル・関数をテストケースとして認識して実行する． それぞれに対して上記のような結果が得られるか実行すると，以下のように表示される．

hogehoge@hogehoge ~ % pytest
======================================================================================================================== test session starts ========================================================================================================================
platform darwin -- Python 3.11.7, pytest-8.0.1, pluggy-1.4.0
plugins: anyio-4.2.0
collected 1 items 

test/test_operations.py .

========================================================================================================================= 1 passed in 3.51s =========================================================================================================================

このようにpassedになっていれば，正しく実行されている． ただassertを書くだけでテストが実行できるのであれば結構便利だと思う．

試しにSelf-Attentionのコードもテストしてみた．

class MultiHeadAttention(nn.Module):
    def __init__(self, dim: int, nhead: int) -> None:
        super(MultiHeadAttention, self).__init__()
        self.W_q = nn.Parameter(torch.randn(nhead, dim, dim // nhead))
        self.W_k = nn.Parameter(torch.randn(nhead, dim, dim // nhead))
        self.W_v = nn.Parameter(torch.randn(nhead, dim, dim // nhead))
        self.softmax = nn.Softmax(dim=-1)
        self.scaler = 1 / math.sqrt(dim)
        self.linear = nn.Linear(dim, dim)
        self.nhead = nhead

    def forward(self, x_: Tensor) -> Tensor:
        B, L, D = x_.shape

        x = x_.repeat(self.nhead, 1, 1, 1).permute(1, 0, 2, 3)

        q = torch.einsum("bhld,hdk->bhlk", x, self.W_q)
        k = torch.einsum("bhld,hdk->bhlk", x, self.W_k)
        v = torch.einsum("bhld,hdk->bhlk", x, self.W_v)

        attn_weight = self.softmax(torch.einsum("bhld,bhkd->bhlk", q, k) * self.scaler)
        qkv = torch.einsum("bhlk,bhkd->bhld", attn_weight, v)

        out = self.linear(qkv.view(B, L, D))

        return out

テストコードをtest_mha.pyとして

from MultiHeadAttention import MultiHeadAttention

def test_MultiHeadAttention():
    mla = MultiHeadAttention(64, 8)
    assert mla(torch.randn(16, 32, 64)).shape == (16, 32, 64)

これを実行すると，

hogehoge@hogehoge ~ % pytest
======================================================================================================================== test session starts ========================================================================================================================
platform darwin -- Python 3.11.7, pytest-8.0.1, pluggy-1.4.0
plugins: anyio-4.2.0
collected 1 items 

test/test_mha.py .

========================================================================================================================= 1 passed in 3.51s =========================================================================================================================

上手く行ってる．

おわりに

今回はpytestを使ったテストコードの書き方をまとめてみた．今後はこれを使って研究開発していきたい