44 lines
1.4 KiB
Python
44 lines
1.4 KiB
Python
#####################################################
|
|
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
|
|
#####################################################
|
|
# pytest ./tests/test_super_vit.py -s #
|
|
#####################################################
|
|
import unittest
|
|
from parameterized import parameterized
|
|
|
|
import torch
|
|
from xautodl.xmodels import transformers
|
|
from xautodl.utils.flop_benchmark import count_parameters
|
|
|
|
|
|
class TestSuperViT(unittest.TestCase):
|
|
"""Test the super re-arrange layer."""
|
|
|
|
def test_super_vit(self):
|
|
model = transformers.get_transformer("vit-base-16")
|
|
tensor = torch.rand((2, 3, 224, 224))
|
|
print("The tensor shape: {:}".format(tensor.shape))
|
|
# print(model)
|
|
outs = model(tensor)
|
|
print("The output tensor shape: {:}".format(outs.shape))
|
|
|
|
@parameterized.expand(
|
|
[
|
|
["vit-cifar10-p4-d4-h4-c32", 32],
|
|
["vit-base-16", 224],
|
|
["vit-large-16", 224],
|
|
["vit-huge-14", 224],
|
|
]
|
|
)
|
|
def test_imagenet(self, name, resolution):
|
|
tensor = torch.rand((2, 3, resolution, resolution))
|
|
config = transformers.name2config[name]
|
|
model = transformers.get_transformer(config)
|
|
outs = model(tensor)
|
|
size = count_parameters(model, "mb", True)
|
|
print(
|
|
"{:10s} : size={:.2f}MB, out-shape: {:}".format(
|
|
name, size, tuple(outs.shape)
|
|
)
|
|
)
|