Distributed training with CPU's

lesscomfortable: <p>Hi!</p>
<p>I am interested in possibly using Ignite to enable distributed training in CPU’s (since I am training a shallow network and have no GPU"s available). I tried using <code>ignite.distributed</code> with the <code>gloo</code> backend, but when I set <code>nproc_per_node</code> to more than 1, the program gets stuck and doesn’t run (it does without setting <code>nproc_per_node</code>). The code is practically the same as the CIFAR <a href="https://github.com/pytorch/ignite/blob/master/examples/contrib/cifar10/main.py" rel="noopener nofollow ugc">example</a>.</p>
<p>My question is, is it possible to train PyTorch models on multiple CPU’s be it with or without Ignite?</p>
<p>Thanks in advance!</p>

vfdev-5: <p><a class="mention" href="/u/lesscomfortable">@lesscomfortable</a> can you share the code and the command you execute to reproduce the issue.</p>
<p>Here is how it would run CIFAR10 script on CPU multi-core (single node) in distributed way:</p>
<pre><code class="lang-bash">CUDA_VISIBLE_DEVICES="" python -m torch.distributed.launch --nproc_per_node=4 --use_env main.py run --backend=gloo
</code></pre>
<p>To ensure that it is not a visual effect that program gets stuck as a single epoch on cifar10 on CPU can several minutes to execute. You can add the progress bar:</p>
<pre><code class="lang-python">    trainer = Engine(train_step)

    if idist.get_rank() == 0 and (not config["with_clearml"]):
        common.ProgressBar(persist=False).attach(trainer)
</code></pre>
<p>In my case it looks like:</p>
<pre><code class="lang-auto">Epoch [1/24]: [4/97]   4%|██████▏                              [00:07&lt;03:45]
</code></pre>
<p>HTH</p>

lesscomfortable: <p>Yes, here is the code. It is almost identical to that in the CIFAR example, except that I’m giving the dataloaders and model to the function myself, I’m using a different loss function and I added a progress bar for training.</p>
<pre><code class="lang-auto">def training(local_rank, config):

    rank = idist.get_rank()
    manual_seed(config["seed"] + rank)
    device = idist.device()

    logger = setup_logger(name="NN-Training")

    log_basic_info(logger, config)

    output_path = config["output_path"]
    if rank == 0:
        if config["stop_iteration"] is None:
            now = datetime.now().strftime("%Y%m%d-%H%M%S")
        else:
            now = f"stop-on-{config['stop_iteration']}"

        folder_name = f"{config['model']}_backend-{idist.backend()}-{idist.get_world_size()}_{now}"
        output_path = Path(output_path) / folder_name
        if not output_path.exists():
            output_path.mkdir(parents=True)
        config["output_path"] = output_path.as_posix()
        logger.info(f"Output path: {config['output_path']}")

        if "cuda" in device.type:
            config["cuda device name"] = torch.cuda.get_device_name(local_rank)

    # Setup dataflow, model, optimizer, criterion
    train_loader, test_loader = config['dataloaders']

    config["num_iters_per_epoch"] = len(train_loader)
    model, optimizer, criterion, lr_scheduler = initialize(config)

    # Create trainer for current task
    trainer = create_trainer(model, optimizer, criterion, lr_scheduler, train_loader.sampler, config, logger)
    
    if idist.get_rank() == 0 and (not config["with_clearml"]):
        common.ProgressBar(desc=f"Trainer (train)", persist=False).attach(trainer, output_transform=lambda x: {'batch loss': x['batch loss']})

    # Let's now setup evaluator engine to perform model's validation and compute metrics
    metrics = {
        "MeanAbsoluteError": MeanAbsoluteError(),
        "Loss": Loss(criterion),
    }

    # We define two evaluators as they wont have exactly similar roles:
    # - `evaluator` will save the best model based on validation score
    evaluator = create_evaluator(model, metrics=metrics, config=config)
    train_evaluator = create_evaluator(model, metrics=metrics, config=config)

    def run_validation(engine):
        epoch = trainer.state.epoch
        state = train_evaluator.run(train_loader)
        log_metrics(logger, epoch, state.times["COMPLETED"], "Train", state.metrics)
        state = evaluator.run(test_loader)
        log_metrics(logger, epoch, state.times["COMPLETED"], "Test", state.metrics)

    trainer.add_event_handler(Events.EPOCH_COMPLETED(every=config["validate_every"]) | Events.COMPLETED, run_validation)

    if rank == 0:
        # Setup TensorBoard logging on trainer and evaluators. Logged values are:
        #  - Training metrics, e.g. running average loss values
        #  - Learning rate
        #  - Evaluation train/test metrics
        evaluators = {"training": train_evaluator, "test": evaluator}
        tb_logger = common.setup_tb_logging(output_path, trainer, optimizer, evaluators=evaluators)

    # Store 2 best models by validation mae starting from num_epochs / 2:
    best_model_handler = Checkpoint(
        {"model": model},
        get_save_handler(config),
        filename_prefix="best",
        n_saved=2,
        global_step_transform=global_step_from_engine(trainer),
        score_name="test_mae",
        score_function=Checkpoint.get_default_score_fn("MeanAbsoluteError"),
    )
    evaluator.add_event_handler(
        Events.COMPLETED(lambda *_: trainer.state.epoch &gt; config["num_epochs"] // 2), best_model_handler
    )

    # In order to check training resuming we can stop training on a given iteration
    if config["stop_iteration"] is not None:

        @trainer.on(Events.ITERATION_STARTED(once=config["stop_iteration"]))
        def _():
            logger.info(f"Stop training on {trainer.state.iteration} iteration")
            trainer.terminate()

    try:
        trainer.run(train_loader, max_epochs=config["num_epochs"])
    except Exception as e:
        logger.exception("")
        raise e

    if rank == 0:
        tb_logger.close()


def run(
    model,
    dataloaders,
    seed=543,
    data_path="/tmp/nn",
    output_path="/tmp/output-nn/",
    batch_size=512,
    momentum=0.9,
    weight_decay=1e-4,
    num_workers=12,
    num_epochs=24,
    learning_rate=0.4,
    num_warmup_epochs=4,
    validate_every=1,
    checkpoint_every=1000,
    backend=None,
    resume_from=None,
    log_every_iters=15,
    nproc_per_node=None,
    stop_iteration=None,
    with_clearml=False,
    with_amp=False,
    **spawn_kwargs,
):
    """Main entry to train an model on NN dataset.
    Args:
        seed (int): random state seed to set. Default, 543.
        data_path (str): input dataset path. Default, "/tmp/nn".
        output_path (str): output path. Default, "/tmp/output-nn".
        model (str): model name (from torchvision) to setup model to train. Default, "resnet18".
        batch_size (int): total batch size. Default, 512.
        momentum (float): optimizer's momentum. Default, 0.9.
        weight_decay (float): weight decay. Default, 1e-4.
        num_workers (int): number of workers in the data loader. Default, 12.
        num_epochs (int): number of epochs to train the model. Default, 24.
        learning_rate (float): peak of piecewise linear learning rate scheduler. Default, 0.4.
        num_warmup_epochs (int): number of warm-up epochs before learning rate decay. Default, 4.
        validate_every (int): run model's validation every ``validate_every`` epochs. Default, 3.
        checkpoint_every (int): store training checkpoint every ``checkpoint_every`` iterations. Default, 1000.
        backend (str, optional): backend to use for distributed configuration. Possible values: None, "nccl", "xla-tpu",
            "gloo" etc. Default, None.
        nproc_per_node (int, optional): optional argument to setup number of processes per node. It is useful,
            when main python process is spawning training as child processes.
        resume_from (str, optional): path to checkpoint to use to resume the training from. Default, None.
        log_every_iters (int): argument to log batch loss every ``log_every_iters`` iterations.
            It can be 0 to disable it. Default, 15.
        stop_iteration (int, optional): iteration to stop the training. Can be used to check resume from checkpoint.
        with_clearml (bool): if True, experiment ClearML logger is setup. Default, False.
        with_amp (bool): if True, enables native automatic mixed precision. Default, False.
        **spawn_kwargs: Other kwargs to spawn run in child processes: master_addr, master_port, node_rank, nnodes
    """
    # catch all local parameters
    config = locals()
    config['dataloaders'] = dataloaders
    config.update(config["spawn_kwargs"])
    del config["spawn_kwargs"]

    spawn_kwargs["nproc_per_node"] = nproc_per_node
    if backend == "xla-tpu" and with_amp:
        raise RuntimeError("The value of with_amp should be False if backend is xla")

    with idist.Parallel(backend=backend, **spawn_kwargs) as parallel:
        parallel.run(training, config)


def get_dataflow(config):
    # - Get train/test datasets
    if idist.get_local_rank() &gt; 0:
        # Ensure that only local rank 0 download the dataset
        # Thus each node will download a copy of the dataset
        idist.barrier()

    train_dataset, test_dataset = utils.get_train_test_datasets(config["data_path"])

    if idist.get_local_rank() == 0:
        # Ensure that only local rank 0 download the dataset
        idist.barrier()

    # Setup data loader also adapted to distributed config: nccl, gloo, xla-tpu
    train_loader = idist.auto_dataloader(
        train_dataset, batch_size=config["batch_size"], num_workers=config["num_workers"], shuffle=True, drop_last=True,
    )

    test_loader = idist.auto_dataloader(
        test_dataset, batch_size=2 * config["batch_size"], num_workers=config["num_workers"], shuffle=False,
    )
    return train_loader, test_loader


def initialize(config):
    model = config["model"]
    # Adapt model for distributed settings if configured
    model = idist.auto_model(model)

    optimizer = optim.SGD(
        model.parameters(),
        lr=config["learning_rate"],
        momentum=config["momentum"],
        weight_decay=config["weight_decay"],
        nesterov=True,
    )
    optimizer = idist.auto_optim(optimizer)
    criterion = nn.L1Loss().to(idist.device())

    le = config["num_iters_per_epoch"]
    milestones_values = [
        (0, 0.0),
        (le * config["num_warmup_epochs"], config["learning_rate"]),
        (le * config["num_epochs"], 0.0),
    ]
    lr_scheduler = PiecewiseLinear(optimizer, param_name="lr", milestones_values=milestones_values)

    return model, optimizer, criterion, lr_scheduler


def log_metrics(logger, epoch, elapsed, tag, metrics):
    metrics_output = "\n".join([f"\t{k}: {v}" for k, v in metrics.items()])
    logger.info(f"\nEpoch {epoch} - Evaluation time (seconds): {elapsed:.2f} - {tag} metrics:\n {metrics_output}")


def log_basic_info(logger, config):
    logger.info(f"Train {config['model']}")
    logger.info(f"- PyTorch version: {torch.__version__}")
    logger.info(f"- Ignite version: {ignite.__version__}")
    if torch.cuda.is_available():
        # explicitly import cudnn as
        # torch.backends.cudnn can not be pickled with hvd spawning procs
        from torch.backends import cudnn

        logger.info(f"- GPU Device: {torch.cuda.get_device_name(idist.get_local_rank())}")
        logger.info(f"- CUDA version: {torch.version.cuda}")
        logger.info(f"- CUDNN version: {cudnn.version()}")

    logger.info("\n")
    logger.info("Configuration:")
    for key, value in config.items():
        logger.info(f"\t{key}: {value}")
    logger.info("\n")

    if idist.get_world_size() &gt; 1:
        logger.info("\nDistributed setting:")
        logger.info(f"\tbackend: {idist.backend()}")
        logger.info(f"\tworld size: {idist.get_world_size()}")
        logger.info("\n")


def create_trainer(model, optimizer, criterion, lr_scheduler, train_sampler, config, logger):

    device = idist.device()

    # Setup Ignite trainer:
    # - let's define training step
    # - add other common handlers:
    #    - TerminateOnNan,
    #    - handler to setup learning rate scheduling,
    #    - ModelCheckpoint
    #    - RunningAverage` on `train_step` output
    #    - Two progress bars on epochs and optionally on iterations

    with_amp = config["with_amp"]
    scaler = GradScaler(enabled=with_amp)

    def train_step(engine, batch):

        x, y = batch[0], batch[1]

        if x.device != device:
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)

        model.train()

        with autocast(enabled=with_amp):
            y_pred = model(x)
            loss = criterion(y_pred, y)

        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        return {
            "batch loss": loss.item(),
        }

    trainer = Engine(train_step)
    trainer.logger = logger

    to_save = {"trainer": trainer, "model": model, "optimizer": optimizer, "lr_scheduler": lr_scheduler}
    metric_names = [
        "batch loss",
    ]

    common.setup_common_training_handlers(
        trainer=trainer,
        train_sampler=train_sampler,
        to_save=to_save,
        save_every_iters=config["checkpoint_every"],
        save_handler=get_save_handler(config),
        lr_scheduler=lr_scheduler,
        output_names=metric_names if config["log_every_iters"] &gt; 0 else None,
        with_pbars=False,
        clear_cuda_cache=False,
    )

    resume_from = config["resume_from"]
    if resume_from is not None:
        checkpoint_fp = Path(resume_from)
        assert checkpoint_fp.exists(), f"Checkpoint '{checkpoint_fp.as_posix()}' is not found"
        logger.info(f"Resume from a checkpoint: {checkpoint_fp.as_posix()}")
        checkpoint = torch.load(checkpoint_fp.as_posix(), map_location="cpu")
        Checkpoint.load_objects(to_load=to_save, checkpoint=checkpoint)

    return trainer


def create_evaluator(model, metrics, config, tag="val"):
    with_amp = config["with_amp"]
    device = idist.device()

    @torch.no_grad()
    def evaluate_step(engine: Engine, batch):
        model.eval()
        x, y = batch[0], batch[1]
        if x.device != device:
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)

        with autocast(enabled=with_amp):
            output = model(x)
        return output, y

    evaluator = Engine(evaluate_step)

    for name, metric in metrics.items():
        metric.attach(evaluator, name)

    if idist.get_rank() == 0 and (not config["with_clearml"]):
        common.ProgressBar(desc=f"Evaluation ({tag})", persist=False).attach(evaluator)

    return evaluator


def get_save_handler(config):
    if config["with_clearml"]:
        from ignite.contrib.handlers.clearml_logger import ClearMLSaver

        return ClearMLSaver(dirname=config["output_path"])

    return DiskSaver(config["output_path"], require_empty=False)
</code></pre>
<p>One notable difference with your command is that I am running this in a Jupyter Notebook in the following manner (no <code>torch.distributed.launch</code> in my code):</p>
<pre><code class="lang-auto">res = run(
    model, 
    [train_dl, test_dl], 
    backend='gloo',
    nproc_per_node=4,
       )
</code></pre>
<p>The result is that the program hangs after printing the following log:</p>
<p><code>ignite.distributed.launcher.Parallel INFO: Spawn function '&lt;function training at ...&gt;' in 4 processes</code></p>
<p>Should I run this as a script with <code>torch.distributed.launch</code>?</p>
<p>Thanks for the prompt answer!</p>

vfdev-5: <p>Here is how to launch the code on Jupyter.</p>
<pre><code class="lang-python">import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""

import time
import ignite.distributed as idist

def training(local_rank, config, **kwargs):
    time.sleep(local_rank)
    print(idist.get_rank(), ': run with config:', config, '- backend=', idist.backend())
    # do the training ...

backend = 'gloo'
dist_configs = {'nproc_per_node': 4, "start_method": "fork"}
config = {'c': 12345}

with idist.Parallel(backend=backend, **dist_configs) as parallel:
    parallel.run(training, config, a=1, b=2)
</code></pre>
<p>You have to use <code>start_method="fork"</code>.</p>
<p>If you would like to run it as a script file and spawn processes from your <code>main.py</code> script as you do, then you can use default <code>start_method</code>. Also, it could be helpful to set <code>persistent_workers=True</code> for the DataLoader to speed up data fetching every epoch…<br>
If you would like to use a script file and spawn processes with <code>torch.distributed.launch</code>, you can simply reuse the command from my previous message (and no need to set <code>persistent_workers=True</code>).</p>

lesscomfortable: <p>Thanks for the tips!</p>
<p>I included all the code necessary to run the code in Jupyter and it now prints out the training info 4, each time with a different rank (1 through 4). However it still hangs, and <code>htop</code> shows that the CPU cores are unused <img src="https://discuss.pytorch.org/images/emoji/apple/grimacing.png?v=10" title=":grimacing:" class="emoji" alt=":grimacing:">. Any other clues on what might be going on?</p>

lesscomfortable: <p>Just in case, these are my parameters:</p>
<pre><code class="lang-auto">|2021-08-09 20:40:58,801 NN-Training INFO: |seed: 543|
|2021-08-09 20:40:58,802 NN-Training INFO: |output_path: test|
|2021-08-09 20:40:58,803 NN-Training INFO: |momentum: 0.9|
|2021-08-09 20:40:58,804 NN-Training INFO: |weight_decay: 0.0001|
|2021-08-09 20:40:58,805 NN-Training INFO: |num_workers: 12|
|2021-08-09 20:40:58,806 NN-Training INFO: |num_epochs: 5|
|2021-08-09 20:40:58,807 NN-Training INFO: |learning_rate: 0.05|
|2021-08-09 20:40:58,809 NN-Training INFO: |num_warmup_epochs: 1|
|2021-08-09 20:40:58,810 NN-Training INFO: |validate_every: 1|
|2021-08-09 20:40:58,811 NN-Training INFO: |checkpoint_every: 3|
|2021-08-09 20:40:58,813 NN-Training INFO: |backend: gloo|
|2021-08-09 20:40:58,818 NN-Training INFO: |resume_from: None|
|2021-08-09 20:40:58,819 NN-Training INFO: |log_every_iters: 15|
|2021-08-09 20:40:58,820 NN-Training INFO: |nproc_per_node: 4|
|2021-08-09 20:40:58,821 NN-Training INFO: |stop_iteration: None|
|2021-08-09 20:40:58,822 NN-Training INFO: |with_clearml: False|
|2021-08-09 20:40:58,823 NN-Training INFO: |with_amp: False|
|2021-08-09 20:40:58,824 NN-Training INFO: |start_method: fork|
</code></pre>

vfdev-5: <p><a class="mention" href="/u/lesscomfortable">@lesscomfortable</a> this is strange. OK, can you please provide pytorch and ignite versions and confirm that you can not run this code snippet: <a href="https://discuss.pytorch.org/t/distributed-training-with-cpus/128983/4" class="inline-onebox">Distributed training with CPU's - #4 by vfdev-5</a> ?</p>
<p>I created a colab to show that above code snippet is running correctly on CPUs: <a href="https://colab.research.google.com/drive/1ss9JUyKYJJuRu5hRjOAbcnGH8ochPVuy" class="inline-onebox" rel="noopener nofollow ugc">Google Colaboratory</a></p>
<p>If you could modify it to bring your code and execute it there (with some random data) and reproduce the issue, it would be nice.</p>

lesscomfortable: <p>This is something I should probably have included in my params: I’m using pytorch 1.6.0. However, your example runs in Colab and my environment successfully when using pytorch 1.6.0 (the problem must be in some other part of my code). These are my versions:</p>
<pre><code class="lang-auto">2021-08-10 09:36:14,511 NN-Training INFO: - PyTorch version: 1.6.0
2021-08-10 09:36:14,511 NN-Training INFO: - Ignite version: 0.4.6
</code></pre>
<p>I’ll try with my code in the Google Colab and let you know if I could reproduce the problem.</p>
<p>Thanks!</p>

vfdev-5: <p>Sounds good! Let me know if you can create a reproducible colab with the issue.</p>
<p>By the way, inspecting the current code, I wonder how do you create data loaders. The way it is done it is unlikely that the training would do a DDP. Data will be the same for all processes and this is probably undesired.</p>

lesscomfortable: <p>After reading a little bit, I am creating them like this (I changed this from the original code I sent previously in this thread):</p>
<pre><code class="lang-auto">train_sampler, valid_sampler = DistributedSampler(train_ds), DistributedSampler(valid_ds)

train_loader = idist.auto_dataloader(
        train_ds, batch_size=bs, num_workers=n_workers, sampler=train_sampler, drop_last=True
)

test_loader = idist.auto_dataloader(
       valid_ds, batch_size=bs, num_workers=n_workers, sampler=valid_sampler
)
</code></pre>

vfdev-5: <p>If you use <code>idist.auto_dataloader</code> you do not need to specify distributed sampler if using the default one.</p>
<p>However, dataloader definition should be inside the training function, such that distributed group is initialized and rank, world size are defined.</p>

lesscomfortable: <p>Ok, I finally made it work! I was trying to reproduce the code in Colab taking your suggestions into account and it now works correctly (I verified that all cpu cores are being used fully). I don’t know exactly what made the trick but here is the working code:</p>
<pre><code class="lang-auto">
def training(local_rank, config, model, datasets):
        
    rank = idist.get_rank()

    print(f'Running with rank {rank} and local rank {rank}')

    manual_seed(config["seed"] + rank)
    device = idist.device()

    logger = setup_logger(name="NN-Training")

    log_basic_info(logger, config)

    output_path = config["output_path"]
    if rank == 0:
        if config["stop_iteration"] is None:
            now = datetime.now().strftime("%Y%m%d-%H%M%S")
        else:
            now = f"stop-on-{config['stop_iteration']}"

        folder_name = f"{config['model']}_backend-{idist.backend()}-{idist.get_world_size()}_{now}"
        output_path = Path(output_path) / folder_name
        if not output_path.exists():
            output_path.mkdir(parents=True)
        config["output_path"] = output_path.as_posix()
        logger.info(f"Output path: {config['output_path']}")

        if "cuda" in device.type:
            config["cuda device name"] = torch.cuda.get_device_name(local_rank)

    train_ds, test_ds = datasets
    
    # Setup data loader also adapted to distributed config: nccl, gloo, xla-tpu
    train_loader = idist.auto_dataloader(
        train_ds, 
        collate_fn=collate,
        batch_size=bs, 
        num_workers=n_workers, 
        drop_last=True
    )

    test_loader = idist.auto_dataloader(
       valid_ds, 
       collate_fn=collate,
       batch_size=bs, 
       num_workers=n_workers, 
    )

    config["num_iters_per_epoch"] = len(train_loader)
    model, optimizer, criterion, lr_scheduler = initialize(model, config)

    # Create trainer for current task
    trainer = create_trainer(model, optimizer, criterion, lr_scheduler, train_loader.sampler, config, logger)

    # Let's now setup evaluator engine to perform model's validation and compute metrics
    metrics = {
        "mae": MeanAbsoluteError()
    }
    
    # Setup running average metric
    acc_metric = RunningAverage(MeanAbsoluteError(output_transform=lambda x: [x['y'], x['y_pred']]), alpha=0.98)
    acc_metric.attach(trainer, 'running_avg_mae')
    
    @trainer.on(Events.EPOCH_COMPLETED(every=1) | Events.COMPLETED)
    def log_running_avg_metrics(engine):
        epoch = trainer.state.epoch
        avg_mae = engine.state.metrics['running_avg_mae']
        logger.info(f"\nEpoch {epoch} - Train metrics:\n running average mae: {avg_mae}")
    
    # - `evaluator` will save the best model based on validation score
    evaluator = create_evaluator(model, metrics=metrics, config=config)

    @trainer.on(Events.EPOCH_COMPLETED(every=config["validate_every"]) | Events.COMPLETED)
    def run_validation(engine):
        epoch = trainer.state.epoch
        state = evaluator.run(test_loader)
        log_metrics(logger, epoch, state.times["COMPLETED"], "Test", state.metrics)

    if rank == 0:
        # Setup TensorBoard logging on trainer and evaluators. Logged values are:
        #  - Training metrics, e.g. running average loss values
        #  - Learning rate
        #  - Evaluation test metrics
        evaluators = {"test": evaluator}
        tb_logger = common.setup_tb_logging(output_path, trainer, optimizer, evaluators=evaluators)

    # Store 2 best models by validation accuracy starting from num_epochs / 2:
    best_model_handler = Checkpoint(
        {"model": model},
        get_save_handler(config),
        filename_prefix="best",
        n_saved=2,
        global_step_transform=global_step_from_engine(trainer),
        score_name="test_mae",
        score_function=Checkpoint.get_default_score_fn("mae"),
    )
    evaluator.add_event_handler(
        Events.COMPLETED(lambda *_: trainer.state.epoch &gt; config["num_epochs"] // 2), best_model_handler
    )

    # In order to check training resuming we can stop training on a given iteration
    if config["stop_iteration"] is not None:

        @trainer.on(Events.ITERATION_STARTED(once=config["stop_iteration"]))
        def _():
            logger.info(f"Stop training on {trainer.state.iteration} iteration")
            trainer.terminate()

    try:
        trainer.run(train_loader, max_epochs=config["num_epochs"])
    except Exception as e:
        logger.exception("")
        raise e

    if rank == 0:
        tb_logger.close()


def run(
    model,
    datasets,
    seed=543,
    output_path="results/",
    momentum=0.9,
    weight_decay=1e-4,
    num_workers=12,
    num_epochs=5,
    learning_rate=0.05,
    num_warmup_epochs=4,
    validate_every=1,
    checkpoint_every=3,
    backend='gloo',
    resume_from=None,
    log_every_iters=15,
    nproc_per_node=None,
    stop_iteration=None,
    with_clearml=False,
    with_amp=False,
    **spawn_kwargs,
):
    # catch all local parameters
    config = locals()
    config.update(config["spawn_kwargs"])
    del config["spawn_kwargs"]

    spawn_kwargs["nproc_per_node"] = nproc_per_node
    if backend == "xla-tpu" and with_amp:
        raise RuntimeError("The value of with_amp should be False if backend is xla")

    with idist.Parallel(backend=backend, **spawn_kwargs) as parallel:
        parallel.run(training, config, model, datasets)

def initialize(model, config):
    # Adapt model for distributed settings if configured
    model = idist.auto_model(model)

    optimizer = optim.SGD(
        model.parameters(),
        lr=config["learning_rate"],
        momentum=config["momentum"],
        weight_decay=config["weight_decay"],
        nesterov=True,
    )
    optimizer = idist.auto_optim(optimizer)
    criterion = nn.L1Loss().to(idist.device())

    le = config["num_iters_per_epoch"]
    milestones_values = [
        (0, 0.0),
        (le * config["num_warmup_epochs"], config["learning_rate"]),
        (le * config["num_epochs"], 0.0),
    ]
    
    lr_scheduler = PiecewiseLinear(optimizer, param_name="lr", milestones_values=milestones_values)

    return model, optimizer, criterion, lr_scheduler


def log_metrics(logger, epoch, elapsed, tag, metrics):
    metrics_output = "\n".join([f"{k}: {v}" for k, v in metrics.items()])
    logger.info(f"\nEpoch {epoch} - {tag} metrics:\n {metrics_output}")

def log_basic_info(logger, config):
    logger.info(f"Train {config['model']} on CIFAR10")
    logger.info(f"- PyTorch version: {torch.__version__}")
    logger.info(f"- Ignite version: {ignite.__version__}")
    if torch.cuda.is_available():
        # explicitly import cudnn as
        # torch.backends.cudnn can not be pickled with hvd spawning procs
        from torch.backends import cudnn

        logger.info(f"- GPU Device: {torch.cuda.get_device_name(idist.get_local_rank())}")
        logger.info(f"- CUDA version: {torch.version.cuda}")
        logger.info(f"- CUDNN version: {cudnn.version()}")

    logger.info("\n")
    logger.info("Configuration:")
    for key, value in config.items():
        logger.info(f"\t{key}: {value}")
    logger.info("\n")

    if idist.get_world_size() &gt; 1:
        logger.info("\nDistributed setting:")
        logger.info(f"\tbackend: {idist.backend()}")
        logger.info(f"\tworld size: {idist.get_world_size()}")
        logger.info("\n")

def create_trainer(model, optimizer, criterion, lr_scheduler, train_sampler, config, logger, tag='train'):

    device = idist.device()

    # Setup Ignite trainer:
    # - let's define training step
    # - add other common handlers:
    #    - TerminateOnNan,
    #    - handler to setup learning rate scheduling,
    #    - ModelCheckpoint
    #    - RunningAverage` on `train_step` output
    #    - Two progress bars on epochs and optionally on iterations

    with_amp = config["with_amp"]
    scaler = GradScaler(enabled=with_amp)

    def train_step(engine, batch):

        x, y = batch[0], batch[1]
        
        if x.device != device:
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)

        model.train()

        with autocast(enabled=with_amp):
            y_pred = model(x)
            loss = criterion(y_pred, y)

        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        return {
            "batch loss": loss.item(),
            "y": y,
            "y_pred": y_pred
        }

    trainer = Engine(train_step)
    trainer.logger = logger

    to_save = {"trainer": trainer, "model": model, "optimizer": optimizer, "lr_scheduler": lr_scheduler}
    metric_names = [
        "batch loss",
    ]

    common.setup_common_training_handlers(
        trainer=trainer,
        train_sampler=train_sampler,
        to_save=to_save,
        save_every_iters=config["checkpoint_every"],
        save_handler=get_save_handler(config),
        lr_scheduler=lr_scheduler,
        output_names=metric_names if config["log_every_iters"] &gt; 0 else None,
        with_pbars=False,
        clear_cuda_cache=False,
    )
    
    if idist.get_rank() == 0 and (not config["with_clearml"]):
        common.ProgressBar(desc=f"Trainer ({tag})", persist=False).attach(trainer, output_transform=lambda x: {'batch loss': x['batch loss']})

    resume_from = config["resume_from"]
    if resume_from is not None:
        checkpoint_fp = Path(resume_from)
        assert checkpoint_fp.exists(), f"Checkpoint '{checkpoint_fp.as_posix()}' is not found"
        logger.info(f"Resume from a checkpoint: {checkpoint_fp.as_posix()}")
        checkpoint = torch.load(checkpoint_fp.as_posix(), map_location="cpu")
        Checkpoint.load_objects(to_load=to_save, checkpoint=checkpoint)

    return trainer

def create_evaluator(model, metrics, config, tag="val"):
    with_amp = config["with_amp"]
    device = idist.device()

    @torch.no_grad()
    def evaluate_step(engine: Engine, batch):
        model.eval()
        x, y = batch[0], batch[1]
        if x.device != device:
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)

        with autocast(enabled=with_amp):
            output = model(x)
        return output, y

    evaluator = Engine(evaluate_step)

    for name, metric in metrics.items():
        metric.attach(evaluator, name)

    if idist.get_rank() == 0 and (not config["with_clearml"]):
        common.ProgressBar(desc=f"Evaluation ({tag})", persist=False).attach(evaluator)

    return evaluator

def get_save_handler(config):
    if config["with_clearml"]:
        from ignite.contrib.handlers.clearml_logger import ClearMLSaver

        return ClearMLSaver(dirname=config["output_path"])

    return DiskSaver(config["output_path"], require_empty=False)
</code></pre>

lesscomfortable: <p>The only thing I could not work out was using <code>persistent_workers=True</code> with <code>idist.auto_dataloader</code>:</p>
<p><code>TypeError: __init__() got an unexpected keyword argument 'persistent_workers'</code>)</p>
<p>Is there any way of using persistent workers with <code>auto_dataloader</code>?</p>

vfdev-5: <p>Glad that you could make it work.</p>
<p>Oh, I just realized that <code>persistent_workers</code> was not available for pytorch 1.6.0 and appears in the docs since 1.7.0. As <code>idist.auto_dataloader</code> just pass through the args and pytorch DataLoader does not have them, that’s why there is an issue.</p>

lesscomfortable: <p>Right, makes sense <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Thank you very much for the help, and let me say Ignite is awesome <img src="https://discuss.pytorch.org/images/emoji/apple/fire.png?v=10" title=":fire:" class="emoji" alt=":fire:"></p>

lesscomfortable: <p>One other thing I noticed is that either <code>num_workers</code> and <code>nproc_per_node</code> must be the same or <code>num_workers</code> must set 0 for the code to run. Do you know the reason behind this?</p>
<p>Example in point: the code runs with <code>num_workers=8</code> and <code>nproc_per_node=8</code> and <code>num_workers=0</code> and <code>nproc_per_node=8</code> but not with <code>num_workers=10/9/7/6</code> (any of these) and <code>nproc_per_node=8</code> (<em>DataLoader worker (pid(s) 53785) exited unexpectedly.</em>).</p>
<p>I’m using <code>persistent_workers=True</code> when running with <code>num_workers&gt;0</code> and <code>persistent_workers=False</code> with <code>num_workers=0</code>.</p>

vfdev-5: <p>Interesting. I’ll try on my infrastructure to see. On the one hand, auto_dataloader scales num_workers per process by the number of local processes: <a href="https://pytorch.org/ignite/generated/ignite.distributed.auto.auto_dataloader.html#ignite.distributed.auto.auto_dataloader" class="inline-onebox" rel="noopener nofollow ugc">auto_dataloader — PyTorch-Ignite v0.4.6 Documentation</a></p>
<pre><code class="lang-python">        if "num_workers" in kwargs and kwargs["num_workers"] &gt;= nproc:
            kwargs["num_workers"] = (kwargs["num_workers"] + nproc - 1) // nproc
</code></pre>
<p>You can see the number of workers the logger is reporting. If <code>num_workers &gt; 8</code>, then it is scaled by <code>nproc_per_node</code>. Otherwise, it will use the number itself. Probable, if you set <code>num_workers=1</code>, it should run as well.</p>
<p>On the other hand, each process of 8 ( <code>nproc_per_node</code>) on the machine will create a dataloader which will also spawn workers. Maybe, the system can not spawn that much ?</p>
<p>Sometimes, it can help to increase shm size if running in docker container with the flag <code>--shm-size 16G</code></p>
<p>EDIT: Updated Colab to run resnet18 on cifar10 with 8 procs and 10 workers, seems like running.</p>

lesscomfortable: <p>Interesting, didn’t know about the scaling of <code>num_workers</code>.</p>
<p>I just peeked into your Colab and it is clear it is working that way. Probably what was happening is that, as long as scaling is on (when <code>num_workers</code> is smaller than 8) the system works as <code>num_workers</code> is low but when scaling is off, after passing a certain threshold (in this case 8), it can’t handle that many workers. This threshold (max <code>num_workers</code>) probably varies depending on the instance.</p>
<p>Will be experimenting to find the sweet spot for my specific case. Thanks!</p>