CineMA/cinema/conv_test.py at main · mathpluscode/CineMA · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
"""Tests for the conv module."""

import pytest
import torch

from cinema.conv import (
    ConvLayerNorm,
    ConvMlp,
    ConvNormActBlock,
    ConvResBlock,
    MaskedConvBlock,
    get_conv_norm,
)


@pytest.mark.parametrize("batch_size", [1, 2])
@pytest.mark.parametrize("in_features", [1, 8])
@pytest.mark.parametrize("image_size", [(16, 8), (7, 16), (16, 8, 7)])
@pytest.mark.parametrize("grad_ckpt", [True, False])
def test_conv_mlp(batch_size: int, in_features: int, image_size: tuple[int, ...], grad_ckpt: bool) -> None:
    """Test output shapes."""
    block = ConvMlp(n_dims=len(image_size), in_features=in_features)
    block.set_grad_ckpt(grad_ckpt)
    for m in block.children():
        if hasattr(m, "grad_ckpt"):
            assert m.grad_ckpt == grad_ckpt
    x = torch.rand(batch_size, in_features, *image_size)
    got = block(x)
    assert got.shape == x.shape


@pytest.mark.parametrize("batch_size", [1, 2])
@pytest.mark.parametrize("in_chans", [1, 8])
@pytest.mark.parametrize("image_size", [(16, 8), (7, 16), (16, 8, 7)])
def test_conv_layernorm(batch_size: int, in_chans: int, image_size: tuple[int, ...]) -> None:
    """Test output shapes."""
    block = ConvLayerNorm(in_chans)
    x = torch.rand(batch_size, in_chans, *image_size)
    got = block(x)
    assert got.shape == x.shape


@pytest.mark.parametrize("batch_size", [1, 2])
@pytest.mark.parametrize("in_chans", [1, 8])
@pytest.mark.parametrize("image_size", [(16, 8), (7, 16), (16, 8, 7)])
@pytest.mark.parametrize("norm", ["instance", "layer", "group"])
def test_get_conv_norm(batch_size: int, in_chans: int, image_size: tuple[int, ...], norm: str) -> None:
    """Test output shapes."""
    block = get_conv_norm(n_dims=len(image_size), in_chans=in_chans, norm=norm)
    x = torch.rand(batch_size, in_chans, *image_size)
    got = block(x)
    assert got.shape == x.shape


@pytest.mark.parametrize("batch_size", [1, 2])
@pytest.mark.parametrize("in_chans", [1, 8])
@pytest.mark.parametrize("out_chans", [1, 8])
@pytest.mark.parametrize("image_size", [(16, 8), (7, 16), (16, 8, 7)])
@pytest.mark.parametrize("norm", ["instance", "layer", "group"])
@pytest.mark.parametrize("grad_ckpt", [True, False])
def test_conv_norm_act(
    batch_size: int, in_chans: int, out_chans: int, image_size: tuple[int, ...], norm: str, grad_ckpt: bool
) -> None:
    """Test output shapes."""
    block = ConvNormActBlock(n_dims=len(image_size), in_chans=in_chans, out_chans=out_chans, norm=norm)
    block.set_grad_ckpt(grad_ckpt)
    for m in block.children():
        if hasattr(m, "grad_ckpt"):
            assert m.grad_ckpt == grad_ckpt
    x = torch.rand(batch_size, in_chans, *image_size)
    got = block(x)
    assert got.shape == (batch_size, out_chans, *image_size)


@pytest.mark.parametrize("batch_size", [1, 2])
@pytest.mark.parametrize("in_chans", [1, 8])
@pytest.mark.parametrize("out_chans", [1, 8])
@pytest.mark.parametrize("image_size", [(16, 8), (7, 16), (16, 8, 7)])
@pytest.mark.parametrize("dropout", [0.0, 0.1])
@pytest.mark.parametrize("norm", ["instance", "layer", "group"])
@pytest.mark.parametrize("grad_ckpt", [True, False])
def test_conv_res(
    batch_size: int,
    in_chans: int,
    out_chans: int,
    image_size: tuple[int, ...],
    dropout: float,
    norm: str,
    grad_ckpt: bool,
) -> None:
    """Test output shapes."""
    block = ConvResBlock(n_dims=len(image_size), in_chans=in_chans, out_chans=out_chans, dropout=dropout, norm=norm)
    block.set_grad_ckpt(grad_ckpt)
    for m in block.children():
        if hasattr(m, "grad_ckpt"):
            assert m.grad_ckpt == grad_ckpt
    x = torch.rand(batch_size, in_chans, *image_size)
    got = block(x)
    assert got.shape == (batch_size, out_chans, *image_size)


@pytest.mark.parametrize("batch_size", [1, 2])
@pytest.mark.parametrize("in_chans", [1, 8])
@pytest.mark.parametrize("image_size", [(16, 8), (7, 16), (16, 8, 7)])
@pytest.mark.parametrize("dropout", [0.0, 0.1])
@pytest.mark.parametrize("drop_path", [0.0, 0.1])
@pytest.mark.parametrize("use_mask", [True, False])
@pytest.mark.parametrize("norm", ["instance", "layer", "group"])
@pytest.mark.parametrize("grad_ckpt", [True, False])
def test_masked_conv_block(
    batch_size: int,
    in_chans: int,
    image_size: tuple[int, ...],
    dropout: float,
    drop_path: float,
    use_mask: bool,
    norm: str,
    grad_ckpt: bool,
) -> None:
    """Test output shapes."""
    block = MaskedConvBlock(n_dims=len(image_size), in_chans=in_chans, dropout=dropout, drop_path=drop_path, norm=norm)
    block.set_grad_ckpt(grad_ckpt)
    for m in block.children():
        if hasattr(m, "grad_ckpt"):
            assert m.grad_ckpt == grad_ckpt
    x = torch.rand(batch_size, in_chans, *image_size)
    mask = torch.rand(batch_size, *image_size) > 0.5 if use_mask else None
    got = block(x, mask)
    assert got.shape == x.shape