nni icon indicating copy to clipboard operation
nni copied to clipboard

sub state dict & Autoformer evolution

Open crj1998 opened this issue 2 years ago • 0 comments

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)

crj1998 avatar Aug 08 '22 02:08 crj1998