This is Hogwild!. Speeding up your Neural Network... | by Phillip Wenig | Towards Data Science

I am currently working on my Ph.D. in the field of distributed machine learning. I’ve noticed that parallel processes aren’t used very often in day-to-day machine learning. Thus, I want to give insights into multiprocessing with PyTorch. In this blog post, you will learn about the Hogwild! algorithm, which is used to run stochastic gradient descent (SGD) in a parallel way. I won’t go into detail as the listed papers cover all the necessary mathematics. Rather, I will explain the general idea and give a very simple example in PyTorch.
If you’ve ever trained a neural network by yourself on your machine, you’ve probably noticed the time and computing power consumed while doing so. Especially, when using big data sets for computer vision or other high dimensional training data. A recent article by MIT has shown, that training an AI algorithm can produce a huge amount of carbon (link to article). This illustrates the necessary power needed for AI model training. Hence, parallel processes for reducing the overall training time should be used if it takes too long.
One possible algorithm is Hogwild! which leverages parallel processes on computers to run SGD simultaneously and asynchronously. Asynchronicity is important for reducing unnecessary idle time in which no calculations are done but energy is still consumed.
Using Hogwild!, each participating process in a parallel training session is responsible for one partition of the data, e.g. having 8 parallel processes would split the data set in 8 equal parts whereas each process is assigned to one part. Moreover, on shared memory or a separate server, an initial model is created which can be accessed by all processes.
Once the training starts, each process loads the current state of the model from shared memory and starts reading the first batch of their data partition. As in standard SGD, each process is calculating the gradients for that batch. The gradients are now directly written to the shared model without blocking the other processes. Once written, the new model parameters are loaded and the next batch is used. Due to the missing blocking, the shared model will sometimes receive old gradients which, one may think, could be a disadvantage. The researchers of Hogwild!, however, could show that the training even benefits from the non-blocking fashion.
Using Hogwild! on a distributed computer cluster has a huge communication overhead because, after each batch, the gradients need to be sent over the network. Thus, Hogwild++ was developed which decreases the amount of overhead by organizing the processes (i.e. each computer) in a ring. While training, a token is sent through that ring. The token bears the global model. Every time a token arrives at a node, the difference between the model weights are calculated and used to update the model in a non-trivial fashion. For more details, the paper is linked in the end.
This code example is inspired by this link and gives an example of how to implement the standard Hogwild! algorithm in PyTorch.
First, we implement a simple image classification model with convolutional layers. The model gives back the LogSoftmax which is useful when using NLLLoss during the training.
class Model(nn.Module):    def __init__(self):        super().__init__()        self.conv1 = nn.Conv2d(3, 10, kernel_size=5)        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)        self.conv2_drop = nn.Dropout2d()        self.fc1 = nn.Linear(500, 10)        self.fc2 = nn.Linear(10, 10)    def forward(self, x, **kwargs):        batch_size = x.shape[0]        x = F.relu(F.max_pool2d(self.conv1(x), 2))        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))        x = x.view(batch_size, -1)        x = F.relu(self.fc1(x))        x = F.dropout(x, training=self.training)        x = self.fc2(x)        x = F.log_softmax(x, dim=1)        return x
Moreover, we define a training function, which is taking a model and a data loader as input parameters. Inside of this function, the Adam optimizer is used and the already mentioned NLLLoss. This training function is a standard implementation of a PyTorch program. Each process participating in Hogwild! will call it at the same time.
def train(model, data_loader):    optimizer = optim.Adam(model.parameters())    criterion = nn.NLLLoss()    for data, labels in tqdm.tqdm(data_loader):        optimizer.zero_grad()        loss = criterion(model(data), labels)        loss.backward()                optimizer.step()
The parallel setup happens in the next code section. Here, we define the number of parallel processes, instantiate the model and push it to shared memory with the single method call share_memory. The dataset used is the CIFAR10 dataset, available in the torchvision package.
We loop over the processes and define a data loader for each process. The data loaders hold a distributed sampler which knows the rank of the process and handles the distribution of the data. Hence, each process has its data partition. The multiprocessing package calls the training function within each process and waits, with the join command, for the processes to finish.
import torch.multiprocessing as mpnum_processes = 4model = Model()model.share_memory()dataset = CIFAR10(            "data",            train=True,            download=True,            transform=transforms.Compose([                transforms.ToTensor(),                transforms.Normalize(                    mean=[0.485, 0.456, 0.406],                    std=[0.229, 0.224, 0.225]                )            ])        )processes = []for rank in range(num_processes):    data_loader = DataLoader(        dataset=dataset,        sampler=DistributedSampler(            dataset=dataset,            num_replicas=num_processes,            rank=rank        ),        batch_size=32    )    p = mp.Process(target=train, args=(model, data_loader))    p.start()    processes.append(p)for p in processes:    p.join()
During training, all the processes have access to the shared model but train only on their very own data partition. Thus, we decrease the training time by approximately 4 which is the number of total training processes.
The whole code can be found on my GitHub repository.
Distributing the training of machine learning becomes more important. Especially, with the increase in training data and data complexity. In general, distributed training saves a lot of time and may also reduce the amount of consumed energy using intelligent distribution techniques. The code example shows, that it is surprisingly easy to apply parallel machine learning training with PyTorch.