Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_serializer.py 3.15 KB
Newer Older
Yuge Zhang's avatar
[Retiarii] Rewrite trainer with PyTorch Lightning (#3359)  
Yuge Zhang committed
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 
import json
from pathlib import Path
import re
import sys

import torch
from nni.retiarii import json_dumps, json_loads, blackbox
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST

sys.path.insert(0, Path(__file__).parent.as_posix())

from imported.model import ImportTest


class Foo:
    def __init__(self, a, b=1):
        self.aa = a
        self.bb = [b + 1 for _ in range(1000)]

    def __eq__(self, other):
        return self.aa == other.aa and self.bb == other.bb


def test_blackbox():
    module = blackbox(Foo, 3)
    assert json_loads(json_dumps(module)) == module
    module = blackbox(Foo, b=2, a=1)
    assert json_loads(json_dumps(module)) == module

    module = blackbox(Foo, Foo(1), 5)
    dumped_module = json_dumps(module)
    assert len(dumped_module) > 200  # should not be too longer if the serialization is correct

    module = blackbox(Foo, blackbox(Foo, 1), 5)
    dumped_module = json_dumps(module)
    assert len(dumped_module) < 200  # should not be too longer if the serialization is correct
    assert json_loads(dumped_module) == module


def test_blackbox_module():
    module = ImportTest(3, 0.5)
    assert json_loads(json_dumps(module)) == module


def test_dataset():
    dataset = blackbox(MNIST, root='data/mnist', train=False, download=True)
    dataloader = blackbox(DataLoader, dataset, batch_size=10)

    dumped_ans = {
        "__type__": "torch.utils.data.dataloader.DataLoader",
        "arguments": {
            "batch_size": 10,
            "dataset": {
                "__type__": "torchvision.datasets.mnist.MNIST",
                "arguments": {"root": "data/mnist", "train": False, "download": True}
            }
        }
    }
    assert json_dumps(dataloader) == json_dumps(dumped_ans)
    dataloader = json_loads(json_dumps(dumped_ans))
    assert isinstance(dataloader, DataLoader)

    dataset = blackbox(MNIST, root='data/mnist', train=False, download=True,
                       transform=blackbox(
                           transforms.Compose,
                           [blackbox(transforms.ToTensor), blackbox(transforms.Normalize, (0.1307,), (0.3081,))]
                       ))
    dataloader = blackbox(DataLoader, dataset, batch_size=10)
    x, y = next(iter(json_loads(json_dumps(dataloader))))
    assert x.size() == torch.Size([10, 1, 28, 28])
    assert y.size() == torch.Size([10])

    dataset = blackbox(MNIST, root='data/mnist', train=False, download=True,
                       transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]))
    dataloader = blackbox(DataLoader, dataset, batch_size=10)
    x, y = next(iter(json_loads(json_dumps(dataloader))))
    assert x.size() == torch.Size([10, 1, 28, 28])
    assert y.size() == torch.Size([10])


def test_type():
    assert json_dumps(torch.optim.Adam) == '{"__typename__": "torch.optim.adam.Adam"}'
    assert json_loads('{"__typename__": "torch.optim.adam.Adam"}') == torch.optim.Adam
    assert re.match(r'{"__typename__": "(.*)test_serializer.Foo"}', json_dumps(Foo))


if __name__ == '__main__':
    test_blackbox()
    test_blackbox_module()
    test_dataset()
    test_type()