Code Review Request

vision
1

Hello all,
I am a PyTorch beginner and was working on Image Classification on MNIST dataset. I have the following code which works, but I wanted to have a code review to get someone else’s opinions / insights and ways to improve.

Apologies if this breaks any rules.

import torch
import torchvision
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt
import torch.nn as nn
from torch.optim import SGD
from sklearn.metrics import accuracy_score


mnist_train = torchvision.datasets.MNIST( '/content/', train = True, download = True, transform = ToTensor() )
mnist_test = torchvision.datasets.MNIST( '/content/', train = False, download = True, transform = ToTensor() )
mnist_train_loader = DataLoader( mnist_train )
mnist_test_loader = DataLoader( mnist_test )


loss_function = nn.CrossEntropyLoss()
optimizer = SGD( params = model.parameters(), lr = 0.01 )


## Training Loop
n_epochs = 2
device = torch.device( 'cuda:0' if torch.cuda.is_available() else 'cpu' )

print( device )

for epoch in range( n_epochs ):


  model.train()
  model.to( device )

  for index, data in enumerate( mnist_train_loader ):

    images, labels = data
    images = images.to( device )
    labels = labels.to( device )


    preds = model( images )

    current_loss = loss_function( preds, labels )

    optimizer.zero_grad()

    current_loss.backward()

    optimizer.step()           


## Evaluation 

pred_list = []
label_list = []

with torch.no_grad():
  model.to( 'cpu' )

  for images, labels in mnist_test_loader:

    images = images.to( 'cpu' )
    labels = labels.to( 'cpu' )


    preds = model( images )
    pred_list.append( preds.argmax().unsqueeze_( 0 ) )
    label_list.append( labels )


accuracy = accuracy_score( label_list, pred_list )

print( f'The model\'s accuracy is { accuracy * 100 } % ' )


## Model 
class CNNClassificationModel( nn.Module ):


  def __init__( self ):
    super( CNNClassificationModel, self ).__init__()
    self.conv1 = nn.Conv2d( 1, 4, 5 )
    self.conv2 = nn.Conv2d( 4, 8, 5 )
    self.conv3 = nn.Conv2d( 8, 16, 5 )
    self.fc1 = nn.Linear( 4096, 2048 )
    self.fc2 = nn.Linear( 2048, 1024 )
    self.fc3 = nn.Linear( 1024, 512 )
    self.fc4 = nn.Linear( 512, 256 )
    self.fc5 = nn.Linear( 256, 128 )
    self.fc6 = nn.Linear( 128, 10 )
    self.relu = nn.ReLU()
    self.softmax = nn.Softmax()


  def forward( self, image ):

    x = self.relu( self.conv1( image ) )
    x = self.relu( self.conv2( x ) )
    x = self.relu( self.conv3( x ) )
    x = x.view( -1, 4096 )
    x = self.relu( self.fc1( x ) )
    x = self.relu( self.fc2( x ) )
    x = self.relu( self.fc3( x ) )
    x = self.relu( self.fc4( x ) )
    x = self.relu( self.fc5( x ) )
    x = self.relu( self.fc6( x ) )

    return x

model = CNNClassificationModel()
2

Please see this it might be help.

import torch
import torchvision
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
import torch.nn as nn
import torch.optim as optim
from sklearn.metrics import accuracy_score

# Define the model architecture
class CNNClassificationModel(nn.Module):
    def __init__(self):
        super(CNNClassificationModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 5 * 5, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(nn.functional.relu(self.conv1(x)))
        x = self.pool(nn.functional.relu(self.conv2(x)))
        x = torch.flatten(x, 1)  # Flatten the tensor for fully connected layers
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Load MNIST dataset
mnist_train = torchvision.datasets.MNIST('/content/', train=True, download=True, transform=ToTensor())
mnist_test = torchvision.datasets.MNIST('/content/', train=False, download=True, transform=ToTensor())
mnist_train_loader = DataLoader(mnist_train, batch_size=64, shuffle=True)
mnist_test_loader = DataLoader(mnist_test, batch_size=64, shuffle=False)

# Initialize the model, loss function, and optimizer
model = CNNClassificationModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# Training Loop
n_epochs = 2
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

for epoch in range(n_epochs):
    model.train()
    running_loss = 0.0
    for inputs, labels in mnist_train_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * inputs.size(0)
    epoch_loss = running_loss / len(mnist_train)
    print(f"Epoch [{epoch + 1}/{n_epochs}], Loss: {epoch_loss:.4f}")

# Evaluation
model.eval()
preds, labels = [], []
for inputs, labels_batch in mnist_test_loader:
    inputs, labels_batch = inputs.to(device), labels_batch.to(device)
    outputs = model(inputs)
    _, predicted = torch.max(outputs, 1)
    preds.extend(predicted.cpu().numpy())
    labels.extend(labels_batch.cpu().numpy())

accuracy = accuracy_score(labels, preds)
print(f'The model\'s accuracy is {accuracy * 100:.2f}%')
3

Remove the softmax as it’s never used and also return the raw logits from the last layer without applying relu on them.

4

Thankyou for taking the time to improve my code, your version is indeed better and more concise.

5

Great points. Thankyou.

6

That sounds good you got you answer enjoy.