def main():
    args = parser.parse_args()
    args, logging, writer = utils.parse_args(args)

    logging.info('# Start Re-training #')

    criterion = LOSS_FACTORY[args.task](args, args.loss_scaling)

    if args.model_type == "stochastic":
        model_temp = STOCHASTIC_FACTORY[args.model]
    else:
        raise NotImplementedError("Other models have not been implemented!")
    model = model_temp(args.input_size, args.output_size, args.layers,
                       args.activation, args, True)

    logging.info('## Model created: ##')
    logging.info(model.__repr__())

    logging.info("### Param size = %f MB, Total number of params = %d ###" %
                 utils.count_parameters_in_MB(model, args))

    logging.info('### Loading model to parallel GPUs ###')

    utils.profile(model, args, logging)
    model = utils.model_to_gpus(model, args)

    logging.info('### Preparing schedulers and optimizers ###')
    optimizer = SGLD(model.parameters(),
                     args.learning_rate,
                     norm_sigma=args.weight_decay)

    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, args.epochs)

    logging.info('## Downloading and preparing data ##')
    train_loader, valid_loader = get_train_loaders(args)

    logging.info('## Beginning Training ##')

    train = Trainer(model, criterion, optimizer, scheduler, args)

    best_error, train_time, val_time = train.train_loop(
        train_loader, valid_loader, logging, writer)

    logging.info(
        '## Finished training, the best observed validation error: {}, total training time: {}, total validation time: {} ##'
        .format(best_error, timedelta(seconds=train_time),
                timedelta(seconds=val_time)))

    logging.info('## Beginning Plotting ##')
    del model
    with torch.no_grad():
        args.samples = 100
        args.model_path = args.save
        model = model_temp(args.input_size, args.output_size, args.layers,
                           args.activation, args, False)
        model = utils.model_to_gpus(model, args)
        model.eval()
        plot_regression_uncertainty(model, PLT, train_loader, args)
        logging.info('# Finished #')
def main(wrapper):
    args = wrapper.args
    model = wrapper.model
    save_dir = args.save

    try:
        wrapper.load()
        logging.info('loaded previously saved weights')
    except Exception as e:
        print(e)

    logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
    logging.info('Args: {}'.format(args))

    if args.eval_only:
        assert save_dir is not None

    data_size = 25000
    time_steps = 1

    B = int(args.epochs * data_size / args.batch_size / time_steps)

    searcher = Random_NAS(B, wrapper, args.seed, save_dir)
    logging.info('budget: %d' % (searcher.B))
    if not args.eval_only:
        searcher.run()
        archs = searcher.get_eval_arch(args.randomnas_rounds, args.n_samples)
    else:
        np.random.seed(args.seed + 1)
        archs = searcher.get_eval_arch(2)
    logging.info(archs)
    #arch = ' '.join([str(a) for a in archs[0][0]])
    arch = str(archs[0][0])
    arch = parse_arch_to_darts('cnn', ast.literal_eval(arch), args.space)
    with open(os.path.join(args.save, 'arch_{}'.format(args.task_id)),
              'w') as f:
        f.write(str(arch))

    logging.info(str(arch))
    utils.write_yaml_results(args, args.results_file_arch, str(arch))
    return arch