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 TestExportCase2(unittest.TestCase):
def setUp(self):
model = ModelLinear()
data_np = np.random.random((3, 64)).astype(np.int64)
self.data = paddle.to_tensor(data_np)
self.ofa_model = OFA(model)
self.ofa_model.set_epoch(0)
outs, _ = self.ofa_model(self.data)
self.config = self.ofa_model.current_config
def test_export_model_linear2(self):
config = self.ofa_model._sample_config(
task='expand_ratio', phase=None, sample_type='smallest')
ex_model = self.ofa_model.export(
config, input_shapes=[[3, 64]], input_dtypes=['int64'])
ex_model(self.data)
assert len(self.ofa_model.ofa_layers) == 3
class TestExport(unittest.TestCase):
def setUp(self):
self._init_model()
def _init_model(self):
self.origin_model = ModelOriginLinear()
model = ModelLinear()
self.ofa_model = OFA(model)
def test_ofa(self):
config = self.ofa_model._sample_config(task='expand_ratio', phase=None)
origin_dict = {}
for name, param in self.origin_model.named_parameters():
origin_dict[name] = param.shape
self.ofa_model.export(
config,
input_shapes=[[1, 64]],
input_dtypes=['int64'],
origin_model=self.origin_model)
for name, param in self.origin_model.named_parameters():
if name in config.keys():
if 'expand_ratio' in config[name]:
assert origin_dict[name][-1] == param.shape[-1] * config[
name]['expand_ratio']
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