Function main

beginner_source/hyperparameter_tuning_tutorial.py:484–536

(num_trials=10, max_num_epochs=10, gpus_per_trial=0, cpus_per_trial=2)

# -----------------------

def main(num_trials=10, max_num_epochs=10, gpus_per_trial=0, cpus_per_trial=2):
    print("Starting hyperparameter tuning.")
    ray.init(include_dashboard=False)
    
    data_dir = os.path.abspath("./data")
    load_data(data_dir)  # Pre-download the dataset
    device = "cuda" if torch.cuda.is_available() else "cpu"
    config = {
        "l1": tune.choice([2**i for i in range(9)]),
        "l2": tune.choice([2**i for i in range(9)]),
        "lr": tune.loguniform(1e-4, 1e-1),
        "batch_size": tune.choice([2, 4, 8, 16]),
        "device": device,
    }
    scheduler = ASHAScheduler(
        max_t=max_num_epochs,
        grace_period=1,
        reduction_factor=2,
    )
    
    tuner = tune.Tuner(
        tune.with_resources(
            partial(train_cifar, data_dir=data_dir),
            resources={"cpu": cpus_per_trial, "gpu": gpus_per_trial}
        ),
        tune_config=tune.TuneConfig(
            metric="loss",
            mode="min",
            scheduler=scheduler,
            num_samples=num_trials,
        ),
        param_space=config,
    )
    results = tuner.fit()

    best_result = results.get_best_result("loss", "min")
    print(f"Best trial config: {best_result.config}")
    print(f"Best trial final validation loss: {best_result.metrics['loss']}")
    print(f"Best trial final validation accuracy: {best_result.metrics['accuracy']}")

    best_trained_model = Net(best_result.config["l1"], best_result.config["l2"])
    best_trained_model = best_trained_model.to(device)
    if gpus_per_trial > 1:
        best_trained_model = nn.DataParallel(best_trained_model)

    best_checkpoint = best_result.checkpoint
    with best_checkpoint.as_directory() as checkpoint_dir:
        checkpoint_path = Path(checkpoint_dir) / "checkpoint.pt"
        best_checkpoint_data = torch.load(checkpoint_path)

        best_trained_model.load_state_dict(best_checkpoint_data["net_state_dict"])
        test_acc = test_accuracy(best_trained_model, device, data_dir)
        print(f"Best trial test set accuracy: {test_acc}")


if __name__ == "__main__":