Multiprocessing batches on CPU with custom layer

Dani16: <p>Hi all,</p>
<p>I created Neural Network with a custom layer in Pytorch which needs to run on CPU and it is made in a way that it can only process one batch at the time. To speed up the training, I would like to use multiprocessing to train such model on N batches in parallel (N being the number of cores of my CPU). Once these batches are processed, I would like to backpropagate the loss and keep training the model with the next N batches. Is there any way to do this using torch.multiprocessing? Or should I just use nn.DataParallel for this?</p>
<p>Thanks in advance for your help!</p>

Jamie_Donnelly: <p>Personally I wouldn’t use <code>multiprocessing</code> to parallelise across physical cores and would instead just go full distributed - which is the easier approach.</p>
<p>Here is an example of how you can do this, essentially all you need to do is write your script and model using <code>DistributedDataParallel</code>, ensure that the backend can find the appropriate <code>WORLD_SIZE</code> variable, and then execute the job through MPI.</p>
<p>Here is an example I just wrote and tested:</p>
<pre><code class="lang-auto">import os 
import torch
from torch import nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

if __name__=="__main__":

    world_size = int(os.environ['OMPI_COMM_WORLD_SIZE']) # to be used internally
    os.environ['WORLD_SIZE'] = str(world_size) # for the backend to find 
    rank = int(os.environ['OMPI_COMM_WORLD_RANK'])

    dist.init_process_group(backend='gloo',init_method='env://',rank=rank)

    if rank == 0:
        print("Process group initialised...")

    model = nn.Linear(10,10)

    ddp_model = DDP(model,device_ids=None)

    print(f"Rank {rank} successfully init'd DDP model...")
</code></pre>
<p><code>DDP</code> handles the parellelisation of the model and you can train it as you would locally, i.e., local calls to <code>backward</code> and <code>opt.step</code> and the communication is abstracted away. If you’re doing CPU-only make sure you use an appropriate backend in the call to <code>init_process_group</code> and then when constructing <code>DDP</code> make sure the <code>device_ids</code> is set to <code>None</code>.</p>
<p>Then to run a script like this on 2 CPUs, I use the following:</p>
<pre><code class="lang-auto">#!/bin/bash

source activate {conda_environment}

export MASTER_ADDR='localhost'

export MASTER_PORT=12340

mpirun -np 2 python main.py
</code></pre>
<p>If you’re doing this on a local machine I would recommend not setting <code>-np</code> to the total number of physical cores on the machine.</p>

Dani16: <p>Hi Jamie, thanks for your reply. I am trying to understand what your script does. In the meantime, I have seen this alternative solution. What do you think?</p>
<pre><code class="lang-auto">import multiprocessing as mp

device = torch.device("cpu")
model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

num_epochs = 10
num_workers = 4  # number of CPU cores to use

def train_worker(batch_idx, batch, model, criterion):
    images, labels = batch
    images = images.to(device)
    labels = labels.to(device)

    # Forward pass
    outputs = model(images)
    loss = criterion(outputs, labels)

    # Backward and optimize
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    print('Batch [{}/{}], Loss: {:.4f}'.format(batch_idx, len(train_loader), loss.item()))

with mp.Pool(processes=num_workers) as pool:
    for epoch in range(num_epochs):
        for i, batch in enumerate(train_loader):
            pool.apply_async(train_worker, args=(i, batch, model, criterion))

        pool.close()
        pool.join()

</code></pre>

Jamie_Donnelly: <p>I don’t think your solution will work, not because you’re using <code>multiprocessing</code> but because the same <code>model</code> is being passed to each worker - no communication occurs between workers and so you won’t get a synchronised model. Personally I find working with <code>multiprocessing</code> more trouble than it’s worth.</p>
<p>This is why you should use PyTorch’s <code>DataParallel</code> or <code>DistributedDataParellel</code> functionality as it will synchronise and communicate gradients between the workers. You wrap the model object in one of these and the hard parts are handled for you.</p>
<p>The code I provided before assigns each instance of python a <code>rank</code> from the environment variables (initialised when executed using <code>mpirun</code>) and sets the <code>WORLD_SIZE</code> environment variable (total number of instances (models trained simultaneously) we want). Then, the <code>dist.init_process_group</code> essentially allows for communication between the instances identified by their <code>rank</code>. (<code>gloo</code> is the default backend for CPU communication).</p>
<p>Then the model is wrapped using <code>DDP</code> which essentially means each python instance trains a replica of the model on different segments of the training data, but by communicating gradients across processes, each model is synchronised i.e., parameters updated with the same gradients. However, you also need to set up a <code>DistributedSampler</code> which will ensure that each process trains on the correct partition of the data so there isn’t overlap and redundant computation. The upside is that when you train the model you write more or less the same code as you would when training the model serially because the communication is handled for you.</p>
<p>Then by executing the script using MPI, you’re running <code>-np</code> instances of python.</p>