A Comprehensive Guide to Training Image Classification Models in PyTorch and TensorFlow

Image classification is one of the most important and widely-used applications of deep learning. Whether it‘s identifying objects in self-driving cars, diagnosing diseases from medical scans, or organizing your personal photos, being able to automatically recognize the content of images has countless valuable use cases.

At the core of image classification are convolutional neural networks (CNNs), a type of deep learning model that can learn rich, hierarchical visual features from labeled image data. Training a CNN to recognize different categories of images is a challenging yet immensely rewarding machine learning problem.

In this guide, we‘ll walk through how to build and train an image classification model from scratch using two of the most popular deep learning frameworks – PyTorch and TensorFlow. We‘ll go through the key steps and best practices that you need to know, and provide full code examples along the way. By the end, you‘ll have a strong understanding of what it takes to get an image classifier up and running.

Let‘s start by introducing the two frameworks we‘ll be using. PyTorch and TensorFlow have a lot in common – they are both open source libraries for defining and training neural networks, with extensive documentation and large, active user communities. However, there are some key differences:

PyTorch

– Defines models using dynamic computation graphs that can be defined and changed on the fly
– "Pythonic" and easy to use, debug and extend
– Ideal for research and rapid prototyping
– Growing quickly in popularity, especially in academia

TensorFlow

– Uses static computation graphs that are defined once before training
– Highly scalable to large models and datasets
– Lots of features for production deployment like model serving
– Most widely used in industry

There‘s no right answer to which one is "better" – it depends on your specific needs and preferences. The good news is that the core concepts for building models are very similar between them. In this guide we‘ll show the code for both so you can compare and contrast.

Alright, now let‘s walk through the steps to build an image classification model, using the classic MNIST handwritten digits dataset as our example. The process can be broken down into 4 main parts: preparing the data, defining the model, training the model, and evaluating performance.

1. Preparing the Data

The first step is to load and preprocess the image data to get it ready for training. For MNIST, we‘ll:

1. Download the dataset (pre-split into 60k training and 10k test images)
2. Convert the images from integers to floats between 0 and 1
3. Normalize by subtracting mean and dividing by standard deviation
4. Create PyTorch Dataset and DataLoader (or tf.data.Dataset for TensorFlow)

Here‘s how that looks in PyTorch:

from torchvision import datasets, transforms

# Define transforms
transform = transforms.Compose([
    transforms.ToTensor(),  # Convert to float tensor
    transforms.Normalize((0.1307,), (0.3081,))  # Normalize with dataset mean and std
])

# Download and load MNIST dataset 
train_data = datasets.MNIST(‘data‘, train=True, download=True, transform=transform)
test_data = datasets.MNIST(‘data‘, train=False, transform=transform)

# Create data loaders  
train_loader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=1000, shuffle=False)

And in TensorFlow:

import tensorflow as tf

# Download and load MNIST dataset
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

# Preprocess the data
x_train = x_train.reshape((60000, 28, 28, 1)) / 255.0
x_test = x_test.reshape((10000, 28, 28, 1)) / 255.0

# Create tf.data datasets  
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(64) 

test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))
test_dataset = test_dataset.batch(1000)

The data loading APIs differ a bit between the two frameworks, but the core idea is the same – download the raw data, apply preprocessing transforms, and wrap in a batched dataset object to efficiently feed the data to your model during training.

2. Defining the Model

Next we need to define the architecture of our CNN model. There are endless possibilities here, but for MNIST a simple model with two convolutional layers, max pooling, and two fully-connected layers does the trick:

# PyTorch  
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

# TensorFlow
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, (3,3), activation=‘relu‘, input_shape=(28, 28, 1)),
    tf.keras.layers.Conv2D(64, (3,3), activation=‘relu‘),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Dropout(0.25),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation=‘relu‘), 
    tf.keras.layers.Dropout(0.5),
    tf.keras.layers.Dense(10, activation=‘softmax‘)
])

The PyTorch code uses the familiar object-oriented style where we define a class inheriting from nn.Module with __init__ and forward methods. The TensorFlow code follows the Keras Sequential style where we stack layers together.

The two models are equivalent though, with the same convolutions, activations, pooling, and dropout layers. Feel free to experiment with the architecture and see how it affects performance!

3. Training the Model

With our data ready and model defined, it‘s time to actually train our image classifier. The training process involves:

1. Defining a loss function to measure how well the model is doing
2. Defining an optimizer that will update the model‘s weights to minimize loss
3. Iterating through the training set for some number of epochs, where each iteration:
   • Samples a batch of data
   • Makes predictions on that batch using the model
   • Computes the loss between predictions and true labels
   • Backpropagates gradients of the loss to update model weights

Here‘s the code to train in PyTorch:

# Loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())

# Training loop
num_epochs = 5
for epoch in range(num_epochs):
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

And in TensorFlow:

# Loss function and optimizer 
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy() 
optimizer = tf.keras.optimizers.Adam()

# Training loop  
num_epochs = 5
for epoch in range(num_epochs):
    for batch_idx, (data, target) in enumerate(train_dataset):
        with tf.GradientTape() as tape:
            predictions = model(data)
            loss = loss_fn(target, predictions)
        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

Again, the code is very similar between the two. The main things to point out:

• We‘re using cross entropy loss because this is a classification problem
• Adam is a popular optimizer that works well for most use cases
• The training loop structure of iterating through batches and epochs is basically identical

As the model trains you should see the loss decreasing over time. You can also track metrics like accuracy to see how performance is improving.

4. Evaluating Performance

Once our model is trained, we need to check how well it actually performs on new data. This is where our test set comes in. We‘ll run the model on the test images and compare its predictions to the true labels to compute evaluation metrics.

The most common metric for classification is simple accuracy – what proportion of images did we predict the correct label for? We can also look at metrics like precision and recall if we want more detail on the types of errors the model is making.

Here‘s the code to evaluate test accuracy in PyTorch:

model.eval()
correct = 0
with torch.no_grad():
    for data, target in test_loader:
        output = model(data)
        pred = output.argmax(dim=1, keepdim=True)  
        correct += pred.eq(target.view_as(pred)).sum().item()

accuracy = correct / len(test_loader.dataset)
print(f‘Test accuracy: {accuracy:.2f}‘)  

And in TensorFlow:

test_loss, test_acc = model.evaluate(test_dataset)
print(f‘Test accuracy: {test_acc:.2f}‘)

With our simple CNN trained for a few epochs, we should expect to see test accuracies in the 97-99% range for MNIST, which is quite good! State-of-the-art models squeeze out a bit more performance with more complex architectures.

Tips and Best Practices

We‘ve covered the core workflow for training an image classifier, but there are some additional techniques that are very useful to know:

• Data augmentation: Artificially increase the size of the training set by applying random transformations like rotations, flips, zooms to images. Helps prevent overfitting.
• Transfer learning: Don‘t train a model from scratch, but instead fine-tune a pre-trained model to your dataset. Much faster and easier, works very well if pre-trained model used similar data.
• k-fold cross validation: More robust way to evaluate model performance. Partition data into k subsets, train and evaluate k times with each subset used once for testing.
• Hyperparameter tuning: The model architecture and training parameters we use are called hyperparameters. Find the best combo of hyperparameters by methodically searching different values (e.g. grid search).

PyTorch or TensorFlow?

So which framework should you use for your own projects? As we‘ve seen, the actual model code is quite similar between PyTorch and TensorFlow these days. The APIs have mostly converged to a simple, object-oriented approach.

PyTorch has some slick features that make it the go-to for many researchers – dynamic computation graphs, simple Pythonic code, easy-to-use higher level APIs. TensorFlow 2 has largely caught up to PyTorch in user experience while retaining its advantages in scalability and production – the ecosystem of tools like TFX, TFLite, TF-Serving is far more mature.

Therefore, which one you pick mostly depends on your use case and environment. Research and rapid experimentation? PyTorch may be the way to go. Building a large-scale production model? TensorFlow is battle-tested.

The good news – the most important deep learning concepts, model architectures, and training techniques will transfer across both. Master one framework and you‘ll have no trouble picking up the other.

Conclusion

We‘ve covered a lot of ground in this guide to training image classifiers in PyTorch and TensorFlow. We walked through all the key steps:

1. Preparing your image data for model training
2. Defining a CNN model architecture
3. Training the model on the data
4. Evaluating model performance on a test set

With the MNIST digits example, we showed the code to implement each part in both frameworks, and discussed additional tips and best practices to take your models to the next level.

At the end of the day, PyTorch and TensorFlow are both powerful, fully-featured tools for building state-of-the-art deep learning models. You can‘t go wrong with either one. The most important thing is to deeply understand the underlying concepts and best practices. Mastering the art of training CNNs will serve you well no matter the framework.

Now, it‘s time for you to go out and train some models! Find an image dataset that interests you and try to build the best classifier you can. Remember that deep learning is an iterative process – starting simple and incrementally growing the complexity of your model and training is the best way to get a feel for it.

Happy training!