def ofa_net(net_id, pretrained=True):
if net_id == 'ofa_proxyless_d234_e346_k357_w1.3':
net = OFAProxylessNASNets(
dropout_rate=0, width_mult=1.3, ks_list=[3, 5, 7], expand_ratio_list=[3, 4, 6], depth_list=[2, 3, 4],
)
elif net_id == 'ofa_mbv3_d234_e346_k357_w1.0':
net = OFAMobileNetV3(
dropout_rate=0, width_mult=1.0, ks_list=[3, 5, 7], expand_ratio_list=[3, 4, 6], depth_list=[2, 3, 4],
)
elif net_id == 'ofa_mbv3_d234_e346_k357_w1.2':
net = OFAMobileNetV3(
dropout_rate=0, width_mult=1.2, ks_list=[3, 5, 7], expand_ratio_list=[3, 4, 6], depth_list=[2, 3, 4],
)
elif net_id == 'ofa_resnet50':
net = OFAResNets(
dropout_rate=0, depth_list=[0, 1, 2], expand_ratio_list=[0.2, 0.25, 0.35], width_mult_list=[0.65, 0.8, 1.0]
)
net_id = 'ofa_resnet50_d=0+1+2_e=0.2+0.25+0.35_w=0.65+0.8+1.0'
else:
raise ValueError('Not supported: %s' % net_id)
if pretrained:
url_base = 'https://hanlab.mit.edu/files/OnceForAll/ofa_nets/'
init = torch.load(
download_url(url_base + net_id, model_dir='.torch/ofa_nets'),
map_location='cpu')['state_dict']
net.load_state_dict(init)
return net
def get_run_manager_and_args():
run_config, args = get_run_config_and_args()
net = OFAMobileNetV3(n_classes=run_config.data_provider.n_classes,
bn_param=(args.bn_momentum, args.bn_eps),
dropout_rate=args.dropout,
base_stage_width=args.base_stage_width,
width_mult=args.width_mult_list,
ks_list=args.ks_list,
expand_ratio_list=args.expand_list,
depth_list=args.depth_list)
if args.kd_ratio > 0:
args.teacher_model = MobileNetV3Large(
n_classes=run_config.data_provider.n_classes,
bn_param=(args.bn_momentum, args.bn_eps),
dropout_rate=0,
width_mult=1.0,
ks=7,
expand_ratio=6,
depth_param=4,
)
distributed_run_manager = RunManager(args.path, net, run_config)
distributed_run_manager.save_config()
return distributed_run_manager, args
else:
# build net from args
args.width_mult_list = [
float(width_mult) for width_mult in args.width_mult_list.split(",")
]
args.ks_list = [int(ks) for ks in args.ks_list.split(",")]
args.expand_list = [int(e) for e in args.expand_list.split(",")]
args.depth_list = [int(d) for d in args.depth_list.split(",")]
args.width_mult_list = (args.width_mult_list[0] if len(
args.width_mult_list) == 1 else args.width_mult_list)
net = OFAMobileNetV3(
n_classes=run_config.data_provider.n_classes,
bn_param=(args.bn_momentum, args.bn_eps),
dropout_rate=args.dropout,
base_stage_width=args.base_stage_width,
width_mult=args.width_mult_list,
ks_list=args.ks_list,
expand_ratio_list=args.expand_list,
depth_list=args.depth_list,
)
# teacher model
if args.kd_ratio > 0:
args.teacher_model = MobileNetV3Large(
n_classes=run_config.data_provider.n_classes,
bn_param=(args.bn_momentum, args.bn_eps),
dropout_rate=0,
width_mult=1.0,
ks=7,
expand_ratio=6,
depth_param=4,
)