test_models.py 3.01 KB
Newer Older
1
2
from collections import OrderedDict
from itertools import product
3
4
5
6
7
import torch
from torchvision import models
import unittest


8
def get_available_classification_models():
9
    # TODO add a registration mechanism to torchvision.models
10
11
12
13
14
15
    return [k for k, v in models.__dict__.items() if callable(v) and k[0].lower() == k[0] and k[0] != "_"]


def get_available_segmentation_models():
    # TODO add a registration mechanism to torchvision.models
    return [k for k, v in models.segmentation.__dict__.items() if callable(v) and k[0].lower() == k[0] and k[0] != "_"]
16
17
18


class Tester(unittest.TestCase):
19
    def _test_classification_model(self, name, input_shape):
20
21
        # passing num_class equal to a number other than 1000 helps in making the test
        # more enforcing in nature
22
        model = models.__dict__[name](num_classes=50)
23
24
25
        model.eval()
        x = torch.rand(input_shape)
        out = model(x)
26
        self.assertEqual(out.shape[-1], 50)
27

28
29
30
31
32
33
34
35
36
37
    def _test_segmentation_model(self, name):
        # passing num_class equal to a number other than 1000 helps in making the test
        # more enforcing in nature
        model = models.segmentation.__dict__[name](num_classes=50, pretrained_backbone=False)
        model.eval()
        input_shape = (1, 3, 300, 300)
        x = torch.rand(input_shape)
        out = model(x)
        self.assertEqual(tuple(out["out"].shape), (1, 50, 300, 300))

38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
    def _make_sliced_model(self, model, stop_layer):
        layers = OrderedDict()
        for name, layer in model.named_children():
            layers[name] = layer
            if name == stop_layer:
                break
        new_model = torch.nn.Sequential(layers)
        return new_model

    def test_resnet_dilation(self):
        # TODO improve tests to also check that each layer has the right dimensionality
        for i in product([False, True], [False, True], [False, True]):
            model = models.__dict__["resnet50"](replace_stride_with_dilation=i)
            model = self._make_sliced_model(model, stop_layer="layer4")
            model.eval()
            x = torch.rand(1, 3, 224, 224)
            out = model(x)
            f = 2 ** sum(i)
            self.assertEqual(out.shape, (1, 2048, 7 * f, 7 * f))

58

59
for model_name in get_available_classification_models():
60
61
62
63
64
65
    # for-loop bodies don't define scopes, so we have to save the variables
    # we want to close over in some way
    def do_test(self, model_name=model_name):
        input_shape = (1, 3, 224, 224)
        if model_name in ['inception_v3']:
            input_shape = (1, 3, 299, 299)
66
67
68
69
70
71
72
73
74
75
        self._test_classification_model(model_name, input_shape)

    setattr(Tester, "test_" + model_name, do_test)


for model_name in get_available_segmentation_models():
    # for-loop bodies don't define scopes, so we have to save the variables
    # we want to close over in some way
    def do_test(self, model_name=model_name):
        self._test_segmentation_model(model_name)
76
77
78
79
80
81

    setattr(Tester, "test_" + model_name, do_test)


if __name__ == '__main__':
    unittest.main()