single_process.py

# from torch.utils.tensorboard import SummaryWriter

import ray
from ray import tune
from ray.tune import CLIReporter
from ray.tune.schedulers import ASHAScheduler

# writer = SummaryWriter()
# import threading
# l = threading.Lock()
def train_function(config):
    print("======================Entering train func===================================")
    # writer.close()
    # l.acquire()
    # l.release()
    for i in range(10):
        tune.report(loss=0, accuracy=1)
    print("======================Leaving train func===================================")

def main(num_samples=10, max_num_epochs=10):
    config = {
        "lr": 0.001,
    }
    scheduler = ASHAScheduler(
        metric="loss",
        mode="min",
        max_t=max_num_epochs,
        grace_period=1,
        reduction_factor=2)
    reporter = CLIReporter(
        # parameter_columns=["l1", "l2", "lr", "batch_size"],
        metric_columns=["loss", "accuracy", "training_iteration"])
    result = tune.run(
        train_function,
        resources_per_trial={"cpu": 1, "gpu": 1},
        config=config,
        num_samples=num_samples,
        scheduler=scheduler,
        progress_reporter=reporter)


if __name__ == "__main__":
    import logging
    ray.init(local_mode=True, num_cpus=1, num_gpus=0, logging_level=logging.DEBUG)
    main()