class TestInputDict(unittest.TestCase):
def setUp(self):
model = ModelInputDict()
sp_net_config = supernet(expand_ratio=[0.5, 1.0])
self.model = Convert(sp_net_config).convert(model)
self.images = paddle.randn(shape=[2, 3, 32, 32], dtype='float32')
self.images2 = {
'data': paddle.randn(shape=[2, 12, 32, 32], dtype='float32')
}
default_run_config = {'skip_layers': ['conv1.0', 'conv2.0']}
self.run_config = RunConfig(**default_run_config)
self.ofa_model = OFA(self.model, run_config=self.run_config)
self.ofa_model._clear_search_space(self.images, data=self.images2)
def test_export(self):
config = self.ofa_model._sample_config(task="expand_ratio",
sample_type="smallest")
self.ofa_model.export(config,
input_shapes=[[1, 3, 32, 32], {
'data': [1, 12, 32, 32]
}],
input_dtypes=['float32', 'float32'])
class TestShortcutSkiplayers(unittest.TestCase):
def setUp(self):
model = ModelShortcut()
sp_net_config = supernet(expand_ratio=[0.5, 1.0])
self.model = Convert(sp_net_config).convert(model)
self.images = paddle.randn(shape=[2, 3, 32, 32], dtype='float32')
self.init_config()
self.ofa_model = OFA(self.model, run_config=self.run_config)
self.ofa_model._clear_search_space(self.images)
def init_config(self):
default_run_config = {'skip_layers': ['branch1.6']}
self.run_config = RunConfig(**default_run_config)
def test_shortcut(self):
self.ofa_model.set_epoch(0)
self.ofa_model.set_task('expand_ratio')
for i in range(5):
self.ofa_model(self.images)
assert list(self.ofa_model._ofa_layers.keys()) == ['branch2.0']
def __call__(self, model, param_state_dict):
paddleslim = try_import('paddleslim')
from paddleslim.nas.ofa import OFA, RunConfig, utils
from paddleslim.nas.ofa.convert_super import Convert, supernet
task = self.ofa_config['task']
expand_ratio = self.ofa_config['expand_ratio']
skip_neck = self.ofa_config['skip_neck']
skip_head = self.ofa_config['skip_head']
run_config = self.ofa_config['RunConfig']
if 'skip_layers' in run_config:
skip_layers = run_config['skip_layers']
else:
skip_layers = []
# supernet config
sp_config = supernet(expand_ratio=expand_ratio)
# convert to supernet
model = Convert(sp_config).convert(model)
skip_names = []
if skip_neck:
skip_names.append('neck.')
if skip_head:
skip_names.append('head.')
for name, sublayer in model.named_sublayers():
for n in skip_names:
if n in name:
skip_layers.append(name)
run_config['skip_layers'] = skip_layers
run_config = RunConfig(**run_config)
# build ofa model
ofa_model = OFA(model, run_config=run_config)
ofa_model.set_epoch(0)
ofa_model.set_task(task)
input_spec = [{
"image": paddle.ones(
shape=[1, 3, 640, 640], dtype='float32'),
"im_shape": paddle.full(
[1, 2], 640, dtype='float32'),
"scale_factor": paddle.ones(
shape=[1, 2], dtype='float32')
}]
ofa_model._clear_search_space(input_spec=input_spec)
ofa_model._build_ss = True
check_ss = ofa_model._sample_config('expand_ratio', phase=None)
# tokenize the search space
ofa_model.tokenize()
# check token map, search cands and search space
logger.info('Token map is {}'.format(ofa_model.token_map))
logger.info('Search candidates is {}'.format(ofa_model.search_cands))
logger.info('The length of search_space is {}, search_space is {}'.
format(len(ofa_model._ofa_layers), ofa_model._ofa_layers))
# set model state dict into ofa model
utils.set_state_dict(ofa_model.model, param_state_dict)
return ofa_model
