xautodl/tests/test_super_norm.py
2021-06-08 23:47:52 -07:00

80 lines
2.9 KiB
Python

#####################################################
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
#####################################################
# pytest ./tests/test_super_norm.py -s #
#####################################################
import unittest
import torch
from xautodl.xlayers import super_core
from xautodl import spaces
class TestSuperSimpleNorm(unittest.TestCase):
"""Test the super simple norm."""
def test_super_simple_norm(self):
out_features = spaces.Categorical(12, 24, 36)
bias = spaces.Categorical(True, False)
model = super_core.SuperSequential(
super_core.SuperSimpleNorm(5, 0.5),
super_core.SuperLinear(10, out_features, bias=bias),
)
print("The simple super module is:\n{:}".format(model))
model.apply_verbose(True)
print(model.super_run_type)
self.assertTrue(model[1].bias)
inputs = torch.rand(20, 10)
print("Input shape: {:}".format(inputs.shape))
outputs = model(inputs)
self.assertEqual(tuple(outputs.shape), (20, 36))
abstract_space = model.abstract_search_space
abstract_space.clean_last()
abstract_child = abstract_space.random()
print("The abstract searc space:\n{:}".format(abstract_space))
print("The abstract child program:\n{:}".format(abstract_child))
model.set_super_run_type(super_core.SuperRunMode.Candidate)
model.enable_candidate()
model.apply_candidate(abstract_child)
output_shape = (20, abstract_child["1"]["_out_features"].value)
outputs = model(inputs)
self.assertEqual(tuple(outputs.shape), output_shape)
def test_super_simple_learn_norm(self):
out_features = spaces.Categorical(12, 24, 36)
bias = spaces.Categorical(True, False)
model = super_core.SuperSequential(
super_core.SuperSimpleLearnableNorm(),
super_core.SuperIdentity(),
super_core.SuperLinear(10, out_features, bias=bias),
)
print("The simple super module is:\n{:}".format(model))
model.apply_verbose(True)
print(model.super_run_type)
self.assertTrue(model[2].bias)
inputs = torch.rand(20, 10)
print("Input shape: {:}".format(inputs.shape))
outputs = model(inputs)
self.assertEqual(tuple(outputs.shape), (20, 36))
abstract_space = model.abstract_search_space
abstract_space.clean_last()
abstract_child = abstract_space.random()
print("The abstract searc space:\n{:}".format(abstract_space))
print("The abstract child program:\n{:}".format(abstract_child))
model.set_super_run_type(super_core.SuperRunMode.Candidate)
model.enable_candidate()
model.apply_candidate(abstract_child)
output_shape = (20, abstract_child["2"]["_out_features"].value)
outputs = model(inputs)
self.assertEqual(tuple(outputs.shape), output_shape)