90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
import paddle
|
|
import paddle.nn as nn
|
|
import paddle.nn.functional as F
|
|
|
|
from ppdet.core.workspace import load_config, merge_config, create
|
|
from ppdet.utils.checkpoint import load_weight, load_pretrain_weight
|
|
from ppdet.utils.logger import setup_logger
|
|
from ppdet.core.workspace import register, serializable
|
|
|
|
from paddle.utils import try_import
|
|
|
|
logger = setup_logger(__name__)
|
|
|
|
|
|
@register
|
|
@serializable
|
|
class OFA(object):
|
|
def __init__(self, ofa_config):
|
|
super(OFA, self).__init__()
|
|
self.ofa_config = ofa_config
|
|
|
|
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
|