diff --git a/test/nas/test_nas.py b/test/nas/test_nas.py index 10673939388..d4a604c1784 100644 --- a/test/nas/test_nas.py +++ b/test/nas/test_nas.py @@ -1,8 +1,6 @@ import os import shutil import unittest -from pathlib import Path - import numpy as np import torch @@ -224,6 +222,118 @@ def test_vision_reference(self): measure_steps=1, ) + def test_parameter_manager_onehot_generic(self): + test_configs = [ + { + 'supernet': 'ofa_mbv3_d234_e346_k357_w1.2', + 'pymoo_vector': [ + 1, 2, 2, 2, 2, 2, 1, 2, 0, 2, 1, 1, 0, 0, 1, + 1, 2, 2, 1, 0, 1, 1, 2, 1, 0, 1, 0, 2, 2, 2, + 0, 0, 2, 2, 2, 2, 1, 1, 2, 1, 2, 0, 2, 0, 0, + ], + 'onehot_vector_expected': [ + 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, + 0, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 1, 0, + 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, + 0, 1, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, + 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, + 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, + 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1, 1, 0, 0, + 0, 0, 1, 1, 0, 0, 1, 0, 0, + ], + }, + { + 'supernet': 'transformer_lt_wmt_en_de', + 'pymoo_vector': [ + 0, 0, 2, 1, 0, 0, 0, 2, 0, 2, 2, 2, 0, 2, 3, 0, 0, 0, 0, 0, + 0, 0, 1, 0, 1, 1, 0, 0, 1, 1, 1, 0, 0, 1, 0, 0, 1, 0, 1], + 'onehot_vector_expected': [ + 1, 0, 1, 0, 0, 0, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0, + 0, 1, 1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 0, 1, + 0, 0, 0, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, + 0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 1, 0, 1, 1, 0, + 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1, 0, + ], + }, + ] + for test_config in test_configs: + nas_agent = NAS('dynas_fake.yaml') + search_algorithm, supernet = 'nsga2', test_config['supernet'] + config = NASConfig(approach='dynas', search_algorithm=search_algorithm) + config.dynas.supernet = supernet + nas_agent = NAS(config) + nas_agent.init_for_search() + + onehot_vector = nas_agent.supernet_manager.onehot_generic(in_array=test_config['pymoo_vector']) + self.assertListEqual(list(onehot_vector), test_config['onehot_vector_expected']) + + def test_parameter_manager_translate2param(self): + test_configs = [ + { + 'supernet': 'ofa_mbv3_d234_e346_k357_w1.2', + 'pymoo_vector': [ + 1, 2, 2, 2, 2, 2, 1, 2, 0, 2, 1, 1, 0, 0, 1, + 1, 2, 2, 1, 0, 1, 1, 2, 1, 0, 1, 0, 2, 2, 2, + 0, 0, 2, 2, 2, 2, 1, 1, 2, 1, 2, 0, 2, 0, 0, + ], + 'param_dict_expected': { + 'd': [4, 2, 4, 2, 2], + 'e': [4, 4, 6, 4, 3, 4, 3, 6, 6, 6, 3, 3, 6, 6, 6, 6, 4, 4, 6, 4], + 'ks': [5, 7, 7, 7, 7, 7, 5, 7, 3, 7, 5, 5, 3, 3, 5, 5, 7, 7, 5, 3], + }, + }, + { + 'supernet': 'transformer_lt_wmt_en_de', + 'pymoo_vector': [0, 0, 2, 1, 0, 0, 0, 2, 0, 2, 2, 2, 0, 2, 3, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 0, 1, 1, 1, 0, 0, 1, 0, 0, 1, 0, 1], + 'param_dict_expected': {'encoder_embed_dim': [640], 'decoder_embed_dim': [640], 'encoder_ffn_embed_dim': [1024, 2048, 3072, 3072, 3072, 1024], 'decoder_ffn_embed_dim': [3072, 1024, 1024, 1024, 3072, 1024], 'decoder_layer_num': [3], 'encoder_self_attention_heads': [8, 8, 8, 8, 8, 8], 'decoder_self_attention_heads': [8, 4, 8, 4, 4, 8], 'decoder_ende_attention_heads': [8, 4, 4, 4, 8, 8], 'decoder_arbitrary_ende_attn': [1, -1, -1, 1, -1, 1]}, + } + ] + + for test_config in test_configs: + nas_agent = NAS('dynas_fake.yaml') + search_algorithm, supernet = 'nsga2', test_config['supernet'] + config = NASConfig(approach='dynas', search_algorithm=search_algorithm) + config.dynas.supernet = supernet + nas_agent = NAS(config) + nas_agent.init_for_search() + + param_dict = nas_agent.supernet_manager.translate2param(test_config['pymoo_vector']) + + self.assertDictEqual(param_dict, test_config['param_dict_expected']) + + + def test_parameter_manager_translate2pymoo(self): + test_configs = [ + { + 'supernet': 'ofa_mbv3_d234_e346_k357_w1.2', + 'param_dict': { + 'd': [4, 2, 4, 2, 2], + 'e': [4, 4, 6, 4, 3, 4, 3, 6, 6, 6, 3, 3, 6, 6, 6, 6, 4, 4, 6, 4], + 'ks': [5, 7, 7, 7, 7, 7, 5, 7, 3, 7, 5, 5, 3, 3, 5, 5, 7, 7, 5, 3], + }, + 'pymoo_vector_expected': [ + 1, 2, 2, 2, 2, 2, 1, 2, 0, 2, 1, 1, 0, 0, 1, + 1, 2, 2, 1, 0, 1, 1, 2, 1, 0, 1, 0, 2, 2, 2, + 0, 0, 2, 2, 2, 2, 1, 1, 2, 1, 2, 0, 2, 0, 0, + ], + }, + { + 'supernet': 'transformer_lt_wmt_en_de', + 'param_dict': {'encoder_embed_dim': [640], 'decoder_embed_dim': [640], 'encoder_ffn_embed_dim': [1024, 2048, 3072, 3072, 3072, 1024], 'decoder_ffn_embed_dim': [3072, 1024, 1024, 1024, 3072, 1024], 'decoder_layer_num': [3], 'encoder_self_attention_heads': [8, 8, 8, 8, 8, 8], 'decoder_self_attention_heads': [8, 4, 8, 4, 4, 8], 'decoder_ende_attention_heads': [8, 4, 4, 4, 8, 8], 'decoder_arbitrary_ende_attn': [1, -1, -1, 1, -1, 1]}, + 'pymoo_vector_expected': [0, 0, 2, 1, 0, 0, 0, 2, 0, 2, 2, 2, 0, 2, 3, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 0, 1, 1, 1, 0, 0, 1, 0, 0, 1, 0, 1], + } + ] + for test_config in test_configs: + nas_agent = NAS('dynas_fake.yaml') + search_algorithm, supernet = 'nsga2', test_config['supernet'] + config = NASConfig(approach='dynas', search_algorithm=search_algorithm) + config.dynas.supernet = supernet + nas_agent = NAS(config) + nas_agent.init_for_search() + + pymoo_vector = nas_agent.supernet_manager.translate2pymoo(test_config['param_dict']) + self.assertListEqual(pymoo_vector, test_config['pymoo_vector_expected']) + if __name__ == "__main__": unittest.main()