autodl-projects/tests/test_super_norm.py

#####################################################
# 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)