Logistic Regression on MNIST with PyTorch | by Asad Mahmood | Towards Data Science

Logistic regression is used to describe data and to explain the relationship between one dependent binary variable and one or more nominal, ordinal, interval or ratio-level independent variables[1]. The figure below shows the difference between Logistic and Linear regression.
In this post, I’ll show how to code a Logistic Regression Model in PyTorch.
We’ll try and solve the classification problem of MNIST dataset. First, let’s import all the libraries we’ll need.
import torchfrom torch.autograd import Variableimport torchvision.transforms as transformsimport torchvision.datasets as dsets
I prefer to keep the following list of steps in front of me when creating a model. This list is present on the PyTorch website [2].
# Step 1. Load Dataset# Step 2. Make Dataset Iterable# Step 3. Create Model Class# Step 4. Instantiate Model Class# Step 5. Instantiate Loss Class# Step 6. Instantiate Optimizer Class# Step 7. Train Model
So let’s go through these steps one by one.
To load the dataset, we make use of torchvision.datasets, a library which has almost all the popular datasets used in Machine Learning. You can check out the complete list of datasets at [3].
train_dataset = dsets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=False)test_dataset = dsets.MNIST(root='./data', train=False, transform=transforms.ToTensor())
We will use the DataLoader class to make our dataset iterable using the following lines of code.
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)
Now, we will create a class that defines the architecture of Logistic Regression.
class LogisticRegression(torch.nn.Module):    def __init__(self, input_dim, output_dim):        super(LogisticRegression, self).__init__()        self.linear = torch.nn.Linear(input_dim, output_dim)    def forward(self, x):        outputs = self.linear(x)        return outputs
Before instantiation, we’ll initialize some parameters like following.
batch_size = 100n_iters = 3000epochs = n_iters / (len(train_dataset) / batch_size)input_dim = 784output_dim = 10lr_rate = 0.001
Now, we initialize our Logistic Regression Model.
model = LogisticRegression(input_dim, output_dim)
We use the cross-entropy to compute the loss.
criterion = torch.nn.CrossEntropyLoss() # computes softmax and then the cross entropy
The optimizer will be the learning algorithm we use. In this case, we will use the Stochastic Gradient Descent.
optimizer = torch.optim.SGD(model.parameters(), lr=lr_rate)
Now in the last step, we’ll train the model using the following code.
iter = 0for epoch in range(int(epochs)):    for i, (images, labels) in enumerate(train_loader):        images = Variable(images.view(-1, 28 * 28))        labels = Variable(labels)        optimizer.zero_grad()        outputs = model(images)        loss = criterion(outputs, labels)        loss.backward()        optimizer.step()        iter+=1        if iter%500==0:            # calculate Accuracy            correct = 0            total = 0            for images, labels in test_loader:                images = Variable(images.view(-1, 28*28))                outputs = model(images)                _, predicted = torch.max(outputs.data, 1)                total+= labels.size(0)                # for gpu, bring the predicted and labels back to cpu fro python operations to work                correct+= (predicted == labels).sum()            accuracy = 100 * correct/total            print("Iteration: {}. Loss: {}. Accuracy: {}.".format(iter, loss.item(), accuracy))
Training, this model for just 3000 iterations gives an accuracy of 82%. You can go ahead and tweak the parameters a bit, to see if the accuracy increases or not.
A good exercise to get a more deep understanding of Logistic Regression models in PyTorch, would be to apply this to any classification problem you could think of. For Example, You could train a Logistic Regression Model to classify the images of your favorite Marvel superheroes (shouldn’t be very hard since half of them are gone :) ).