Momentum for Deep Learning Made Easy (How-To)

Hope that explanation made sense. Stepping through it always helps. If you’re still having a bit of trouble, maybe write it out by hand. That helped me when I was learning momentum.

Let’s take a look at how momentum works in our model training loop.

First thing first, we need some data. We’re going to use PyTorch’s randn function to create a tensor for our X values. The tensor will consist of random values, and will have the shape 1000×1000. We’ll also create a tensor for our y values having the shape 1000×1. We’ll make it binary — half the values will be set 0, while the other half will be set to 1.

dim = 1000
X = torch.randn((dim,dim), dtype=float)
y = torch.cat((torch.zeros(dim//2, dtype=float), torch.ones(dim//2, dtype=float)))

Awesome. Now that we have our data, we’re going to use PyTorch’s Dataset and DataLoader functionality so that we can perform batch gradient descent (if you’re unfamiliar with PyTorch’s Dataset and DataLoader classes, then please check out my article which explains it in detail).

class ExampleDataset(Dataset):    def __init__(self, X, y):
self.X = X
self.y = y
self.len = len(self.X)    def __getitem__(self, index):
return self.X[index], self.y[index]    def __len__(self):
return self.lendata = ExampleDataset(X, y)
data_loader = DataLoader(dataset=data, batch_size=64, shuffle=True)

We’re now ready to create our neural network architecture using PyTorch’s nn.Module class. Our architecture will be very straightforward. It’s a simple feed forward network consisting of 1 hidden layer, which will be half the size of our input layer (i.e., 500). To keep things simple, the activation function we’ll use is the sigmoid function.

class Model(nn.Module):    def __init__(self, dim):
super().__init__()
self.hidden_layer = nn.Linear(dim, dim//2)
self.last_layer = nn.Linear(dim//2, 1)
self.sigmoid = nn.Sigmoid()    def forward(self, X):
out = self.sigmoid(self.hidden_layer(X))
out = self.sigmoid(self.last_layer(out))
return out

For the first round of training, we’re going to train without momentum. It allows us to set a benchmark and compare the performance with and without momentum.

Now, something we’re going to do is use the GPU (i.e., CUDA) for this example. It’s easy to use and it will allow us to train our model faster. All we need to do is place our model and our data onto the GPU by calling the .to function. This will alert PyTorch that we want to use the GPU when training and PyTorch will take care of the rest.

Note: to turn the GPU on in Google Colab, go to Edit > Notebook settings and then chagne the hardware settings to GPU (the GPU is free to use).

We’ll instantiate our Model class, initialize our loss function and optimizer. We’ll train for a total 2000 epochs with a learning rate of 1e-5.

Note: momentum should be done with a small learning rate, otherwise we won’t see optimal results.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
epochs = 2000
learning_rate = 1e-5model = Model(dim).to(device)
bce = nn.BCELoss()
optimizer = torch.optim.SGD(model.parameters(), lr = learning_rate)

Let’s train our model!

epoch_loss = []
model.train()
for epoch in range(epochs+1):
iteration_loss = 0.
for i, data in enumerate(data_loader):
X, y = data
y_pred = model(X.float().to(device))
loss = bce(y_pred.to(device), y.reshape(-1,1).float().to(device))
iteration_loss += loss.item()        optimizer.zero_grad()
loss.backward()
optimizer.step()    epoch_loss.append(iteration_loss/(i+1))
if(epoch % 100 == 0):
print('Train: epoch: {0} - loss: {1:.5f}'.format(epoch, epoch_loss[-1]))

plt.plot(epoch_loss)

model = Model(dim).to(device)
bce = nn.BCELoss()
# we'll reuse the learning_rate variable from above
optimizer = torch.optim.SGD(model.parameters(), lr = learning_rate)
momentum = 0.9

epoch_loss = []
velocities = {}model.train()
for epoch in range(epochs+1):
iteration_loss = 0.
for i, data in enumerate(data_loader):
X, y = data
y_pred = model(X.float().to(device))
loss = bce(y_pred.to(device), y.reshape(-1,1).float().to(device))
iteration_loss += loss.item()        optimizer.zero_grad()
loss.backward()
for name, w in model.named_parameters():
if name not in velocities:
velocities[name] = torch.zeros_like(w.grad).detach()
velocities[name].mul_(momentum).add_(w.grad)
w.grad = velocities[name]
optimizer.step()    epoch_loss.append(iteration_loss/(i+1))
if(epoch % 100 == 0):
print('Train: epoch: {0} - loss: {1:.5f}'.format(epoch, epoch_loss[-1]))

plt.plot(epoch_loss)

model = Model(dim).to(device)
bce = nn.BCELoss()
# we'll reuse the learning_rate and momentum variables from above
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum)

epoch_loss = []
model.train()
for epoch in range(epochs+1):
iteration_loss = 0.
for i, data in enumerate(data_loader):
X, y = data
y_pred = model(X.float().to(device))
loss = bce(y_pred.to(device), y.reshape(-1,1).float().to(device))
iteration_loss += loss.item()        optimizer.zero_grad()
loss.backward()        optimizer.step()    epoch_loss.append(iteration_loss/(i+1))
if(epoch % 100 == 0):
print('Train: epoch: {0} - loss: {1:.5f}'.format(epoch, epoch_loss[-1]))

plt.plot(epoch_loss)

torch.optim.SGD??