nni
nni copied to clipboard
sub state dict & Autoformer evolution
Description
Add fixed subnet state dict load operation and Autoformer evolution
How to test
import nni.retiarii.hub.pytorch as hub
from nni.retiarii.hub.pytorch.utils.fixed import FixedFactory
from nni.retiarii.utils import sub_state_dict
arch = {'embed_dim': 192, 'depth': 12}
mlp_ratio = [3.0] * 14
num_head = [3] * 14
arch.update({f'mlp_ratio_{k}': v for k, v in enumerate(mlp_ratio)})
arch.update({f'num_head_{k}': v for k, v in enumerate(num_head)})
init_kwargs = {
'search_embed_dim': (240, 216, 192),
'search_mlp_ratio': (4.0, 3.5, 3.0),
'search_num_heads': (4, 3),
'search_depth': (14, 13, 12),
'qkv_bias': True,
'drop_rate': 0.0,
'drop_path_rate': 0.1,
'global_pool': True,
'num_classes': 1000
}
model = FixedFactory(hub.AutoformerSpace, arch)(**init_kwargs)
supernet_state_dict = torch.load(CHECKPOINT)
state_dict = sub_state_dict(model, supernet_state_dict )
model.load_state_dict(state_dict)