xautodl/tests/test_super_vit.py

41 lines
1.4 KiB
Python
Raw Normal View History

2021-06-09 11:16:56 +02:00
#####################################################
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
#####################################################
# pytest ./tests/test_super_vit.py -s #
#####################################################
import sys
import unittest
import torch
from xautodl.xmodels import transformers
from xautodl.utils.flop_benchmark import count_parameters
2021-06-09 14:39:35 +02:00
2021-06-09 11:16:56 +02:00
class TestSuperViT(unittest.TestCase):
"""Test the super re-arrange layer."""
def test_super_vit(self):
2021-06-09 14:39:35 +02:00
model = transformers.get_transformer("vit-base-16")
tensor = torch.rand((16, 3, 224, 224))
2021-06-09 11:16:56 +02:00
print("The tensor shape: {:}".format(tensor.shape))
2021-06-09 14:39:35 +02:00
# print(model)
2021-06-09 11:16:56 +02:00
outs = model(tensor)
print("The output tensor shape: {:}".format(outs.shape))
2021-06-09 14:39:35 +02:00
def test_imagenet(self):
2021-06-09 11:16:56 +02:00
name2config = transformers.name2config
2021-06-09 14:39:35 +02:00
print("There are {:} models in total.".format(len(name2config)))
2021-06-09 11:16:56 +02:00
for name, config in name2config.items():
2021-06-09 14:39:35 +02:00
if "cifar" in name:
tensor = torch.rand((16, 3, 32, 32))
else:
tensor = torch.rand((16, 3, 224, 224))
2021-06-09 11:16:56 +02:00
model = transformers.get_transformer(config)
2021-06-09 14:39:35 +02:00
outs = model(tensor)
2021-06-09 11:16:56 +02:00
size = count_parameters(model, "mb", True)
2021-06-09 14:39:35 +02:00
print(
"{:10s} : size={:.2f}MB, out-shape: {:}".format(
name, size, tuple(outs.shape)
)
)