colossalai
62 строки · 1.7 Кб
1import pytest
2import timm.models as tmm
3import torch
4import torchvision.models as tm
5
6from colossalai.fx._compatibility import is_compatible_with_meta
7
8if is_compatible_with_meta():
9from colossalai.fx.profiler import MetaTensor
10
11from colossalai.testing import clear_cache_before_run
12
13tm_models = [
14tm.vgg11,
15tm.resnet18,
16tm.densenet121,
17tm.mobilenet_v3_small,
18tm.resnext50_32x4d,
19tm.wide_resnet50_2,
20tm.regnet_x_16gf,
21tm.mnasnet0_5,
22tm.efficientnet_b0,
23]
24
25tmm_models = [
26tmm.resnest.resnest50d,
27tmm.beit.beit_base_patch16_224,
28tmm.cait.cait_s24_224,
29tmm.efficientnet.efficientnetv2_m,
30tmm.resmlp_12_224,
31tmm.vision_transformer.vit_base_patch16_224,
32tmm.deit_base_distilled_patch16_224,
33tmm.convnext.convnext_base,
34tmm.vgg.vgg11,
35tmm.dpn.dpn68,
36tmm.densenet.densenet121,
37tmm.rexnet.rexnet_100,
38tmm.swin_transformer.swin_base_patch4_window7_224,
39]
40
41
42@pytest.mark.skipif(not is_compatible_with_meta(), reason="torch version is lower than 1.12.0")
43@clear_cache_before_run()
44def test_torchvision_models():
45for m in tm_models:
46model = m()
47data = torch.rand(100000, 3, 224, 224, device="meta")
48model(MetaTensor(data, fake_device=torch.device("cpu"))).sum().backward()
49
50
51@pytest.mark.skipif(not is_compatible_with_meta(), reason="torch version is lower than 1.12.0")
52@clear_cache_before_run()
53def test_timm_models():
54for m in tmm_models:
55model = m()
56data = torch.rand(100000, 3, 224, 224, device="meta")
57model(MetaTensor(data, fake_device=torch.device("cpu"))).sum().backward()
58
59
60if __name__ == "__main__":
61test_torchvision_models()
62test_timm_models()
63