131 lines
5.0 KiB
Python
131 lines
5.0 KiB
Python
#####################################################
|
|
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
|
|
#####################################################
|
|
# pytest ./tests/test_super_model.py -s #
|
|
#####################################################
|
|
import torch
|
|
import unittest
|
|
|
|
from xautodl.xlayers import super_core
|
|
from xautodl import spaces
|
|
|
|
|
|
class TestSuperLinear(unittest.TestCase):
|
|
"""Test the super linear."""
|
|
|
|
def test_super_linear(self):
|
|
out_features = spaces.Categorical(12, 24, 36)
|
|
bias = spaces.Categorical(True, False)
|
|
model = super_core.SuperLinear(10, out_features, bias=bias)
|
|
print("The simple super linear module is:\n{:}".format(model))
|
|
model.apply_verbose(True)
|
|
|
|
print(model.super_run_type)
|
|
self.assertTrue(model.bias)
|
|
|
|
inputs = torch.rand(20, 10)
|
|
print("Input shape: {:}".format(inputs.shape))
|
|
print("Weight shape: {:}".format(model._super_weight.shape))
|
|
print("Bias shape: {:}".format(model._super_bias.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["_out_features"].value)
|
|
outputs = model(inputs)
|
|
self.assertEqual(tuple(outputs.shape), output_shape)
|
|
|
|
def test_super_mlp_v1(self):
|
|
hidden_features = spaces.Categorical(12, 24, 36)
|
|
out_features = spaces.Categorical(24, 36, 48)
|
|
mlp = super_core.SuperMLPv1(10, hidden_features, out_features)
|
|
print(mlp)
|
|
mlp.apply_verbose(False)
|
|
self.assertTrue(mlp.fc1._out_features, mlp.fc2._in_features)
|
|
|
|
inputs = torch.rand(4, 10)
|
|
outputs = mlp(inputs)
|
|
self.assertEqual(tuple(outputs.shape), (4, 48))
|
|
|
|
abstract_space = mlp.abstract_search_space
|
|
print(
|
|
"The abstract search space for SuperMLPv1 is:\n{:}".format(abstract_space)
|
|
)
|
|
self.assertEqual(
|
|
abstract_space["fc1"]["_out_features"],
|
|
abstract_space["fc2"]["_in_features"],
|
|
)
|
|
self.assertTrue(
|
|
abstract_space["fc1"]["_out_features"]
|
|
is abstract_space["fc2"]["_in_features"]
|
|
)
|
|
|
|
abstract_space.clean_last()
|
|
abstract_child = abstract_space.random(reuse_last=True)
|
|
print("The abstract child program is:\n{:}".format(abstract_child))
|
|
self.assertEqual(
|
|
abstract_child["fc1"]["_out_features"].value,
|
|
abstract_child["fc2"]["_in_features"].value,
|
|
)
|
|
|
|
mlp.set_super_run_type(super_core.SuperRunMode.Candidate)
|
|
mlp.enable_candidate()
|
|
mlp.apply_candidate(abstract_child)
|
|
outputs = mlp(inputs)
|
|
output_shape = (4, abstract_child["fc2"]["_out_features"].value)
|
|
self.assertEqual(tuple(outputs.shape), output_shape)
|
|
|
|
def test_super_mlp_v2(self):
|
|
hidden_multiplier = spaces.Categorical(1.0, 2.0, 3.0)
|
|
out_features = spaces.Categorical(24, 36, 48)
|
|
mlp = super_core.SuperMLPv2(10, hidden_multiplier, out_features)
|
|
print(mlp)
|
|
mlp.apply_verbose(False)
|
|
|
|
inputs = torch.rand(4, 10)
|
|
outputs = mlp(inputs)
|
|
self.assertEqual(tuple(outputs.shape), (4, 48))
|
|
|
|
abstract_space = mlp.abstract_search_space
|
|
print(
|
|
"The abstract search space for SuperMLPv2 is:\n{:}".format(abstract_space)
|
|
)
|
|
|
|
abstract_space.clean_last()
|
|
abstract_child = abstract_space.random(reuse_last=True)
|
|
print("The abstract child program is:\n{:}".format(abstract_child))
|
|
|
|
mlp.set_super_run_type(super_core.SuperRunMode.Candidate)
|
|
mlp.enable_candidate()
|
|
mlp.apply_candidate(abstract_child)
|
|
outputs = mlp(inputs)
|
|
output_shape = (4, abstract_child["_out_features"].value)
|
|
self.assertEqual(tuple(outputs.shape), output_shape)
|
|
|
|
def test_super_stem(self):
|
|
out_features = spaces.Categorical(24, 36, 48)
|
|
model = super_core.SuperAlphaEBDv1(6, out_features)
|
|
inputs = torch.rand(4, 360)
|
|
|
|
abstract_space = model.abstract_search_space
|
|
abstract_space.clean_last()
|
|
abstract_child = abstract_space.random(reuse_last=True)
|
|
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)
|
|
outputs = model(inputs)
|
|
output_shape = (4, 60, abstract_child["_embed_dim"].value)
|
|
self.assertEqual(tuple(outputs.shape), output_shape)
|