29 lines
892 B
Python
29 lines
892 B
Python
|
#####################################################
|
||
|
|
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
|
||
|
|
#####################################################
|
||
|
|
# pytest ./tests/test_super_rearrange.py -s #
|
||
|
|
#####################################################
|
||
|
|
import sys
|
||
|
|
import unittest
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
lib_dir = (Path(__file__).parent / "..").resolve()
|
||
|
|
print("LIB-DIR: {:}".format(lib_dir))
|
||
|
|
if str(lib_dir) not in sys.path:
|
||
|
|
sys.path.insert(0, str(lib_dir))
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from xautodl import xlayers
|
||
|
|
|
||
|
|
|
||
|
|
class TestSuperReArrange(unittest.TestCase):
|
||
|
|
"""Test the super re-arrange layer."""
|
||
|
|
|
||
|
|
def test_super_re_arrange(self):
|
||
|
|
layer = xlayers.SuperReArrange(
|
||
|
|
"b c (h p1) (w p2) -> b (h w) (c p1 p2)", p1=4, p2=4
|
||
|
|
)
|
||
|
|
tensor = torch.rand((8, 4, 32, 32))
|
||
|
|
print("The tensor shape: {:}".format(tensor.shape))
|
||
|
|
print(layer)
|