".github/workflows/build-wheels-macos.yml" did not exist on "be3f120798183d8e8d8a420f22f2b2223f1b0701"
test_hub.py 1.7 KB
Newer Older
1
2
3
4
5
import torch.hub as hub
import tempfile
import shutil
import os
import sys
Anirudh's avatar
Anirudh committed
6
import pytest
7
8
9
10
11
12
13
14
15


def sum_of_model_parameters(model):
    s = 0
    for p in model.parameters():
        s += p.sum()
    return s


16
SUM_OF_PRETRAINED_RESNET18_PARAMS = -12703.9931640625
17
18


Anirudh's avatar
Anirudh committed
19
20
21
@pytest.mark.skipif('torchvision' in sys.modules,
                    reason='TestHub must start without torchvision imported')
class TestHub:
22
23
24
25
26
27
28
29
30
31
32
33
    # Only run this check ONCE before all tests start.
    # - If torchvision is imported before all tests start, e.g. we might find _C.so
    #   which doesn't exist in downloaded zip but in the installed wheel.
    # - After the first test is run, torchvision is already in sys.modules due to
    #   Python cache as we run all hub tests in the same python process.

    def test_load_from_github(self):
        hub_model = hub.load(
            'pytorch/vision',
            'resnet18',
            pretrained=True,
            progress=False)
Anirudh's avatar
Anirudh committed
34
        assert sum_of_model_parameters(hub_model).item() == pytest.approx(SUM_OF_PRETRAINED_RESNET18_PARAMS)
35
36
37
38
39
40
41
42
43

    def test_set_dir(self):
        temp_dir = tempfile.gettempdir()
        hub.set_dir(temp_dir)
        hub_model = hub.load(
            'pytorch/vision',
            'resnet18',
            pretrained=True,
            progress=False)
Anirudh's avatar
Anirudh committed
44
45
        assert sum_of_model_parameters(hub_model).item() == pytest.approx(SUM_OF_PRETRAINED_RESNET18_PARAMS)
        assert os.path.exists(temp_dir + '/pytorch_vision_master')
46
47
48
49
        shutil.rmtree(temp_dir + '/pytorch_vision_master')

    def test_list_entrypoints(self):
        entry_lists = hub.list('pytorch/vision', force_reload=True)
Anirudh's avatar
Anirudh committed
50
        assert 'resnet18' in entry_lists
51
52
53


if __name__ == "__main__":
Anirudh's avatar
Anirudh committed
54
    pytest.main([__file__])