Building a Binary Image Classifier with PyTorch: A Step-by-Step Guide

ByJordan Brown

Image classification is a fundamental problem in computer vision, with the goal of automatically assigning a label to an input image from a predefined set of categories. In the case of binary image classification, there are only two possible labels, making it a great starting point for learning how to build image classifiers with deep learning.

Some interesting applications of binary image classifiers include:

  • Medical imaging: Classifying X-rays or MRI scans as normal or abnormal to detect diseases
  • Quality control: Identifying defective products on an assembly line to maintain quality
  • Safety monitoring: Checking if workers are wearing required safety gear like hard hats
  • Content moderation: Detecting inappropriate or explicit content in user-uploaded images

In this guide, we‘ll walk through the process of building and training a binary image classifier using PyTorch – an open-source deep learning framework known for its simplicity and flexibility. By the end, you‘ll have a working classifier that can predict image classes with high accuracy, along with the knowledge to apply these techniques to your own datasets and tasks.

The Dataset

For this tutorial, we‘ll use the Cats vs Dogs dataset, a collection of 25,000 images of cats and dogs split into training and testing sets. The dataset is available from the Kaggle website and is a great resource for getting started with image classification.

Here are a few sample images from the dataset:

[Insert 2-3 sample cat and dog images side by side]

As you can see, the images vary quite a bit in terms of lighting, pose, background, etc. The classifier will need to learn to be invariant to these types of variations to accurately predict the correct animal.

Preprocessing the Data

Before we can train a classifier on the cat and dog images, we need to do some preprocessing to get the data into the right format and improve convergence when training the model. Some standard preprocessing steps for image data include:

  1. Resizing: All images need to be the same dimensions, so we typically resize them to a fixed size like 224×224
  2. Normalization: Pixel values are normalized to a small range like [-1, 1] to make optimization more stable
  3. Data augmentation: The training images are randomly perturbed with transforms like rotations, flips, and crops to increase the effective size of the dataset and help the model generalize better to unseen images

Here‘s how those preprocessing steps can be chained together using PyTorch‘s transforms:

train_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.RandomCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

test_transform = transforms.Compose([
    transforms.Resize(256), 
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

The training and testing transforms are slightly different – we only apply random data augmentation to the training set.

We can then create PyTorch datasets that apply these transforms and load the data in batches using a DataLoader:

train_dataset = datasets.ImageFolder(train_dir, transform=train_transform)
test_dataset = datasets.ImageFolder(test_dir, transform=test_transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64)

Defining the Model Architecture

Next up is defining the architecture of our convolutional neural network (CNN) that will learn to classify cats and dogs. CNNs have been very successful for image-related tasks, as the convolution operation allows the network to learn patterns like edges and textures that are useful for prediction.

We‘ll use a simple CNN architecture with 3 convolutional layers and 2 fully-connected layers:

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3)  # 3 input channels, 32 output channels, 3x3 convolution  
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.conv3 = nn.Conv2d(64, 128, 3)
        self.pool = nn.MaxPool2d(2, 2)  # 2x2 max pooling
        self.fc1 = nn.Linear(128 * 26 * 26, 512)  # 26x26 is the dimension after conv and pool layers
        self.fc2 = nn.Linear(512, 1)  # Binary classifier has 1 output (with sigmoid activation)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x))) 
        x = self.pool(F.relu(self.conv3(x)))
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

The forward method defines the sequence of operations applied to an input batch of images to generate predicted probabilities. It consists of alternating convolution and max pooling layers to learn hierarchical features, followed by flattening and fully-connected layers to make the final prediction.

Training the Model

With the model architecture defined, we can now train it to classify cats and dogs. The training process involves repeatedly showing the model batches of images, comparing its predictions to the true labels, and adjusting the model‘s weights to minimize the prediction error.

To quantify the error, we define a loss function that measures how far off the model‘s predictions are from the target labels. For binary classification, binary cross entropy (BCE) loss is commonly used. And to adjust the model‘s weights to reduce the loss, we use an optimizer like stochastic gradient descent (SGD).

Here‘s the code to set up the loss function and optimizer:

criterion = nn.BCEWithLogitsLoss()  # Combines a Sigmoid layer and the BCELoss
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

Then we can train the model for a number of epochs, where each epoch is a full pass through the training set:

num_epochs = 10

for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        optimizer.zero_grad()

        outputs = net(images)
        loss = criterion(outputs, labels.float().unsqueeze(1))

        loss.backward()
        optimizer.step()

For each batch, we first zero out the gradients from the previous iteration. Then we pass the batch of images through the model to get the predicted probabilities. We calculate the BCE loss between the predictions and true labels. Finally, we backpropagate the gradients and update the model parameters using the optimizer.

It‘s important to track the model‘s performance during training to spot issues like overfitting. We can calculate metrics like accuracy and loss on the training set and a held-out validation set after each epoch:

net.eval()  # Put model in evaluation mode to turn off dropout etc.

train_correct = 0
train_total = 0

with torch.no_grad():  # No need to compute gradients when evaluating
    for images, labels in train_loader:
        outputs = net(images)
        predicted = torch.round(torch.sigmoid(outputs))
        train_total += labels.size(0)
        train_correct += (predicted == labels.unsqueeze(1)).sum().item()

print(f"Training accuracy: {train_correct / train_total}")

If training is progressing well, the training and validation metrics should improve over time until they plateau.

Evaluating the Trained Model

Once the model is done training, we can evaluate its final performance on the test set. In addition to accuracy, it‘s informative to look at a confusion matrix showing a breakdown of correct and incorrect predictions for each class:

test_correct = 0
test_total = 0
true_pos = 0
true_neg = 0 
false_pos = 0
false_neg = 0

with torch.no_grad():
    for images, labels in test_loader:
        outputs = net(images)
        predicted = torch.round(torch.sigmoid(outputs))
        test_total += labels.size(0)
        test_correct += (predicted == labels.unsqueeze(1)).sum().item()

        true_pos += ((predicted == 1) & (labels.unsqueeze(1) == 1)).sum().item()
        true_neg += ((predicted == 0) & (labels.unsqueeze(1) == 0)).sum().item()
        false_pos += ((predicted == 1) & (labels.unsqueeze(1) == 0)).sum().item()
        false_neg += ((predicted == 0) & (labels.unsqueeze(1) == 1)).sum().item()

print(f"Testing accuracy: {test_correct / test_total}")
print(f"Confusion matrix:")
print(f"               Actual Cat   Actual Dog") 
print(f"Predicted Cat    {true_pos:<10}  {false_pos:<10}")
print(f"Predicted Dog    {false_neg:<10}  {true_neg:<10}")

We can also visualize some of the model‘s predictions on test images to get a qualitative sense of its performance:

# Visualize some predictions
fig, axs = plt.subplots(2, 5, figsize=(20, 8))
axs = axs.flatten()

for images, labels in test_loader:
    for i in range(10):  
        with torch.no_grad():
            outputs = net(images[i].unsqueeze(0))
            predicted = torch.round(torch.sigmoid(outputs))

        img = images[i].permute(1, 2, 0)  # Reorder dimensions for plotting
        axs[i].imshow(img.numpy())
        axs[i].set_title(f"Pred: {predicted.item()} | Actual: {labels[i].item()}")
        axs[i].axis(‘off‘)
    break

This will plot 10 test images in a grid with the model‘s predicted label and true label for each.

Tips for Improving Performance

If you‘re not satisfied with your binary image classifier‘s performance, here are some tips for potentially improving it:

  1. Use a more powerful CNN architecture like ResNet, DenseNet, or EfficientNet. These have been carefully designed and are widely used for image classification.

  2. Try transfer learning by starting with a CNN pretrained on a large dataset like ImageNet and fine-tuning it on your specific dataset. The pretrained weights can give the model a helpful starting point.

  3. Gather additional training data or generate more using data augmentation techniques. In general, more data leads to better deep learning models.

  4. Experiment with different learning rates, optimizers, regularization methods, etc. The right combination can help the model converge faster and generalize better.

  5. Ensemble multiple models trained with different architectures or initializations and combine their predictions. Ensembles tend to outperform individual models.

Conclusion

In this guide, we walked through the key steps for building a binary image classifier using PyTorch:

  1. Preprocess the data by resizing, normalizing, and augmenting the images
  2. Define a CNN architecture to extract useful features and make predictions
  3. Train the model on the data using a loss function and optimizer
  4. Evaluate the model‘s performance on a test set with metrics like accuracy and a confusion matrix
  5. Experiment with techniques for improving the model like transfer learning and ensembling

With some patience and iteration, you can train highly accurate image classifiers using the power of deep learning and PyTorch. By understanding the core concepts and code covered here, you‘ll be well-equipped to take on more advanced computer vision tasks and deploy your models to real-world applications.

Some helpful resources for learning more:

How useful was this post?

Click on a star to rate it!

Average rating 0 / 5. Vote count: 0

No votes so far! Be the first to rate this post.

You May Like to Read,

Similar Posts