torch.manual_seed(args.manual_seed)
torch.cuda.manual_seed_all(args.manual_seed)
np.random.seed(args.manual_seed)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
os.makedirs(args.path, exist_ok=True)
# build run config from args
args.lr_schedule_param = None
args.opt_param = {
'momentum': args.momentum,
'nesterov': not args.no_nesterov,
}
run_config = ImagenetRunConfig(**args.__dict__)
print('run_config.save_path', run_config.save_path)
# debug, adjust run_config
if args.debug:
run_config.train_batch_size = 32
run_config.test_batch_size = 32
run_config.valid_size = None
run_config.slice_idx = 20480
width_stages_str = '-'.join(args.width_stages.split(','))
# build net from args
args.width_stages = [int(val) for val in args.width_stages.split(',')]
args.n_cell_stages = [int(val) for val in args.n_cell_stages.split(',')]
args.stride_stages = [int(val) for val in args.stride_stages.split(',')]
args.conv_candidates = [
'3x3_MBConv3',
torch.manual_seed(args.manual_seed)
torch.cuda.manual_seed_all(args.manual_seed)
np.random.seed(args.manual_seed)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
os.makedirs(args.path, exist_ok=True)
# build run config from args
args.lr_schedule_param = None
args.opt_param = {
'momentum': args.momentum,
'nesterov': not args.no_nesterov,
}
if args.dataset == "imagenet":
run_config = ImagenetRunConfig(**args.__dict__)
elif args.dataset == "speech_commands":
run_config = SpeechCommandsRunConfig(**args.__dict__)
else:
raise NotImplementedError
# debug, adjust run_config
if args.debug:
run_config.train_batch_size = 256
run_config.test_batch_size = 256
run_config.valid_size = 256
run_config.n_worker = 0
width_stages_str = '-'.join(args.width_stages.split(','))
# build net from args
args.width_stages = [int(val) for val in args.width_stages.split(',')]
def main(args, myargs):
# args = parser.parse_args()
torch.manual_seed(args.manual_seed)
torch.cuda.manual_seed_all(args.manual_seed)
np.random.seed(args.manual_seed)
# os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
os.makedirs(args.path, exist_ok=True)
# build run config from args
args.lr_schedule_param = None
args.opt_param = {
'momentum': args.momentum,
'nesterov': not args.no_nesterov,
}
run_config = ImagenetRunConfig(
**args.__dict__
)
# debug, adjust run_config
if args.debug:
run_config.train_batch_size = 256
run_config.test_batch_size = 256
run_config.valid_size = 256
run_config.n_worker = 0
width_stages_str = '-'.join(args.width_stages.split(','))
# build net from args
args.width_stages = [int(val) for val in args.width_stages.split(',')]
args.n_cell_stages = [int(val) for val in args.n_cell_stages.split(',')]
args.stride_stages = [int(val) for val in args.stride_stages.split(',')]
args.conv_candidates = [
'3x3_MBConv3', '3x3_MBConv6',
'5x5_MBConv3', '5x5_MBConv6',
'7x7_MBConv3', '7x7_MBConv6',
]
super_net = SuperProxylessNASNets(
width_stages=args.width_stages, n_cell_stages=args.n_cell_stages, stride_stages=args.stride_stages,
conv_candidates=args.conv_candidates, n_classes=run_config.data_provider.n_classes, width_mult=args.width_mult,
bn_param=(args.bn_momentum, args.bn_eps), dropout_rate=args.dropout
)
# build arch search config from args
if args.arch_opt_type == 'adam':
args.arch_opt_param = {
'betas': (args.arch_adam_beta1, args.arch_adam_beta2),
'eps': args.arch_adam_eps,
}
else:
args.arch_opt_param = None
if args.target_hardware is None:
args.ref_value = None
else:
args.ref_value = ref_values[args.target_hardware]['%.2f' % args.width_mult]
if args.arch_algo == 'grad':
from nas_manager import GradientArchSearchConfig
if args.grad_reg_loss_type == 'add#linear':
args.grad_reg_loss_params = {'lambda': args.grad_reg_loss_lambda}
elif args.grad_reg_loss_type == 'mul#log':
args.grad_reg_loss_params = {
'alpha': args.grad_reg_loss_alpha,
'beta': args.grad_reg_loss_beta,
}
else:
args.grad_reg_loss_params = None
arch_search_config = GradientArchSearchConfig(**args.__dict__)
elif args.arch_algo == 'rl':
from nas_manager import RLArchSearchConfig
arch_search_config = RLArchSearchConfig(**args.__dict__)
else:
raise NotImplementedError
print('Run config:')
for k, v in run_config.config.items():
print('\t%s: %s' % (k, v))
print('Architecture Search config:')
for k, v in arch_search_config.config.items():
print('\t%s: %s' % (k, v))
# arch search run manager
arch_search_run_manager = ArchSearchRunManager(args.path, super_net, run_config, arch_search_config)
# resume
if args.resume:
try:
arch_search_run_manager.load_model()
except Exception:
from pathlib import Path
home = str(Path.home())
warmup_path = os.path.join(
home, 'Workspace/Exp/arch_search/%s_ProxylessNAS_%.2f_%s/warmup.pth.tar' %
(run_config.dataset, args.width_mult, width_stages_str)
)
if os.path.exists(warmup_path):
print('load warmup weights')
arch_search_run_manager.load_model(model_fname=warmup_path)
else:
print('fail to load models')
# warmup
if arch_search_run_manager.warmup:
arch_search_run_manager.warm_up(warmup_epochs=args.warmup_epochs)
# joint training
arch_search_run_manager.train(fix_net_weights=args.debug)