%%capture
!uv pip install torch torchvision numpy matplotlib scipy tqdm safetensors

# Importing necessary libraries

import os
from datetime import datetime

import torch  # Core PyTorch library for tensor operations
import torch.nn as nn  # Neural network components
import torch.optim as optim  # Optimization algorithms
import numpy as np  # Numerical operations
import matplotlib.pyplot as plt  # Plotting|
from torchvision import datasets, transforms  # Datasets and transformations for computer vision
from torch.utils.data import DataLoader, Dataset, random_split, Subset  # Data handling utilities
from tqdm import tqdm
from safetensors.torch import save_model, load_model

# Define transformations for the dataset
transform = transforms.ToTensor()

# Load the training and test datasets
trainset = datasets.SVHN(root='./data', split='train', download=True, transform=transform)
testset = datasets.SVHN(root='./data', split='test', download=True, transform=transform)

# Set batch size as a parameter
batch_size = 64


# Function to create data loaders with a specified batch size
def get_dataloader(dataset, batch_size, shuffle):
  return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)


# Create data loaders for training and testing
trainloader = get_dataloader(trainset, batch_size=batch_size, shuffle=True)
testloader = get_dataloader(testset, batch_size=batch_size, shuffle=False)

Using downloaded and verified file: ./data/train_32x32.mat
Using downloaded and verified file: ./data/test_32x32.mat

# Define the class labels for the SVHN dataset
classes = [str(i) for i in range(10)]


# Function to display a grid of images with specified rows, columns, and class labels
def show_images_grid(dataloader, rows, cols):
  # Get a batch of training images and labels
  dataiter = iter(dataloader)
  images, labels = next(dataiter)

  fig, axes = plt.subplots(rows, cols, figsize=(cols * 2, rows * 2))
  idx = 0
  for i in range(rows):
    for j in range(cols):
      if idx < len(images):
        npimg = images[idx].numpy()
        axes[i, j].imshow(np.transpose(npimg, (1, 2, 0)))
        axes[i, j].set_title(f'Class: {classes[labels[idx]]}')
        axes[i, j].axis('off')
        idx += 1
  plt.tight_layout()
  plt.show()

# Display a 2x5 grid of images from trainloader
show_images_grid(trainloader, rows=2, cols=5)

del trainset, testset, trainloader, testloader, transform, batch_size, get_dataloader

# Set up device for training (use GPU if available)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

Using device: cuda

batch_size = 64  # Batch size for data loading
learning_rate = 0.05  # Learning rate for optimizer
num_epochs = 20  # Number of training epochs
momentum = 0.9
weight_decay = 1e-4

# Define transformations for SVHN dataset
transform = transforms.Compose([
  transforms.ToTensor(),
  transforms.Normalize((0.4377, 0.4438, 0.4728), (0.1980, 0.2010, 0.1970)),
])

# Load training and test datasets with transformations
train_dataset = datasets.SVHN(root='./data', split='train', download=True, transform=transform)
test_dataset = datasets.SVHN(root='./data', split='test', download=True, transform=transform)

# Create data loaders with specified batch size
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

Using downloaded and verified file: ./data/train_32x32.mat
Using downloaded and verified file: ./data/test_32x32.mat

# utilities to serde with safetensors


def save_pretrained(model, name='model', *, base_path='./model'):
  os.makedirs(base_path, exist_ok=True)
  timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')

  # Create filename with model name and timestamp
  filename = f'{name}_{timestamp}.safetensors'
  filepath = os.path.join(base_path, filename)

  # Save the model
  save_model(model, filepath)
  print(f'Model saved to {filepath}')
  return filepath

class PretrainedMixin:
  @classmethod
  def from_pretrained(cls, filepath, device='cuda'):
    model = cls().to(device)
    load_model(model, filepath)
    model.eval()
    return model

  def save_pretrained(self, base_path='./model'):
    save_pretrained(self, name=self.__class__.__qualname__, base_path=base_path)

class SVHNClassifier(nn.Module, PretrainedMixin):
  def __init__(self):
    super(SVHNClassifier, self).__init__()

    # not specified in spec, but add dropout for stability
    self.convblock1 = nn.Sequential(
      nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
      nn.BatchNorm2d(32),
      nn.ReLU(),
      nn.MaxPool2d(kernel_size=2, stride=2),
    )

    self.convblock2 = nn.Sequential(
      nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),
      nn.BatchNorm2d(64),
      nn.ReLU(),
      nn.MaxPool2d(kernel_size=2, stride=2),
    )

    self.convblock3 = nn.Sequential(
      nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
      nn.BatchNorm2d(128),
      nn.ReLU(),
      nn.MaxPool2d(kernel_size=2, stride=2),
    )

    # Calculate input size for the first fully connected layer
    # Input image: 32x32
    # After 3 max pooling layers (32 -> 16 -> 8 -> 4)
    # With 128 channels: 128 * 4 * 4 = 2048
    self.fc = nn.Sequential(nn.Linear(128 * 4 * 4, 128), nn.ReLU(), nn.Linear(128, 10))

  def forward(self, x):
    x = self.convblock1(x)
    x = self.convblock2(x)
    x = self.convblock3(x)
    x = x.view(x.size(0), -1)
    x = self.fc(x)
    return x

# Initialize model and move to device
model = SVHNClassifier().to(device)

# Define loss function and optimizer
criterion = nn.CrossEntropyLoss()  # CrossEntropy includes softmax
if True:
  optimizer = optim.SGD(model.parameters(), lr=learning_rate)
if False:
  optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=weight_decay)

# Define function to evaluate model performance
def evaluate(model, data_loader, criterion, epoch, num_epochs, *, ncols=100):
  model.eval()  # Set model to evaluation mode
  total_loss = 0
  correct = 0
  total = 0

  with (
    torch.no_grad(),
    tqdm(data_loader, desc=f'eval  {epoch + 1}/{num_epochs}', ncols=ncols) as eval_data_loader,
  ):  # No gradient calculation for evaluation
    for data, targets in eval_data_loader:
      # Move data to device
      data = data.to(device)
      targets = targets.to(device)

      # Forward pass
      outputs = model(data)
      loss = criterion(outputs, targets)

      # Accumulate loss
      total_loss += loss.item() * data.size(0)

      # Calculate accuracy
      _, predicted = torch.max(outputs.data, 1)
      total += targets.size(0)
      correct += (predicted == targets).sum().item()

  # Calculate average loss and accuracy
  avg_loss = total_loss / total
  accuracy = 100 * correct / total

  return accuracy, avg_loss

# Training the model
train_losses, test_losses = [], []
train_accuracies, test_accuracies = [], []

# bar stuff
ncols = 100

for epoch in range(num_epochs):
  model.train()  # Set model to training mode
  running_loss = 0.0
  total = 0
  correct = 0

  # Training loop with tqdm progress bar
  with tqdm(
    enumerate(train_loader), total=len(train_loader), desc=f'epoch {epoch + 1}/{num_epochs}', ncols=ncols
  ) as batch_pbar:
    for i, (data, targets) in batch_pbar:
      # Move data to device
      data = data.to(device)
      targets = targets.to(device)

      # Zero the gradients
      optimizer.zero_grad()

      # Forward pass
      outputs = model(data)
      loss = criterion(outputs, targets)

      # Backward pass and optimize
      loss.backward()
      optimizer.step()

      # Track training statistics
      running_loss += loss.item() * data.size(0)
      _, predicted = torch.max(outputs.data, 1)
      total += targets.size(0)
      correct += (predicted == targets).sum().item()

      # Update batch progress bar
      batch_pbar.set_postfix({'batch_loss': loss.item()})

  # Calculate average training loss and accuracy for this epoch
  train_loss = running_loss / total
  train_accuracy = 100 * correct / total

  # Evaluate on train and test sets after each epoch
  test_accuracy, test_loss = evaluate(model, test_loader, criterion, epoch + 1, num_epochs, ncols=ncols)

  # Print epoch statistics
  print('-' * 60)
  print(f'epoch {epoch + 1:02d}/{num_epochs:02d} - Train Loss: {train_loss:.6f}, Train Acc: {train_accuracy:.2f}%')
  print(f'            - Test Loss: {test_loss:.6f}, Test Acc: {test_accuracy:.2f}%')
  print('-' * 60)

  # Store losses and accuracies
  train_losses.append(train_loss)
  test_losses.append(test_loss)
  train_accuracies.append(train_accuracy)
  test_accuracies.append(test_accuracy)

epoch 1/20: 100%|█████████████████████████████| 1145/1145 [00:12<00:00, 91.35it/s, batch_loss=0.283]
eval  2/20: 100%|████████████████████████████████████████████████| 407/407 [00:03<00:00, 114.90it/s]

------------------------------------------------------------
epoch 01/20 - Train Loss: 0.738209, Train Acc: 76.31%
            - Test Loss: 0.460860, Test Acc: 86.17%
------------------------------------------------------------

epoch 2/20: 100%|█████████████████████████████| 1145/1145 [00:12<00:00, 91.66it/s, batch_loss=0.448]
eval  3/20: 100%|████████████████████████████████████████████████| 407/407 [00:03<00:00, 112.73it/s]

------------------------------------------------------------
epoch 02/20 - Train Loss: 0.363855, Train Acc: 89.01%
            - Test Loss: 0.348315, Test Acc: 89.72%
------------------------------------------------------------

epoch 3/20: 100%|█████████████████████████████| 1145/1145 [00:12<00:00, 91.11it/s, batch_loss=0.373]
eval  4/20: 100%|████████████████████████████████████████████████| 407/407 [00:03<00:00, 114.52it/s]

------------------------------------------------------------
epoch 03/20 - Train Loss: 0.294560, Train Acc: 91.20%
            - Test Loss: 0.375101, Test Acc: 88.85%
------------------------------------------------------------

epoch 4/20: 100%|█████████████████████████████| 1145/1145 [00:12<00:00, 92.84it/s, batch_loss=0.072]
eval  5/20: 100%|████████████████████████████████████████████████| 407/407 [00:03<00:00, 114.35it/s]

------------------------------------------------------------
epoch 04/20 - Train Loss: 0.249752, Train Acc: 92.66%
            - Test Loss: 0.316894, Test Acc: 90.64%
------------------------------------------------------------

epoch 5/20: 100%|█████████████████████████████| 1145/1145 [00:12<00:00, 93.01it/s, batch_loss=0.282]
eval  6/20: 100%|████████████████████████████████████████████████| 407/407 [00:03<00:00, 112.55it/s]

------------------------------------------------------------
epoch 05/20 - Train Loss: 0.218510, Train Acc: 93.53%
            - Test Loss: 0.290605, Test Acc: 91.79%
------------------------------------------------------------

def plot_metrics(train_metrics, test_metrics, metric_name):
  plt.figure(figsize=(8, 6))
  epochs = np.arange(len(train_metrics))

  plt.plot(epochs, train_metrics, label=f'Train {metric_name}', color='blue')
  plt.plot(epochs, test_metrics, label=f'Test {metric_name}', color='red')

  plt.xlabel('Epochs')
  plt.ylabel(metric_name)
  plt.title(f'{metric_name} over Epochs')
  plt.legend()
  plt.grid(True)
  plt.tight_layout()
  plt.show()

# Plot loss
plot_metrics(train_losses, test_losses, 'Loss')

# Plot accuracy
plot_metrics(train_accuracies, test_accuracies, 'Accuracy')

if False:
  model.save_pretrained()

del (
  device,
  batch_size,
  learning_rate,
  num_epochs,
  train_dataset,
  test_dataset,
  train_loader,
  test_loader,
  transform,
  SVHNClassifier,
  model,
  criterion,
  optimizer,
  evaluate,
  train_losses,
  test_losses,
  train_accuracies,
  test_accuracies,
  show_images_grid,
  classes,
)

# Set up device for training (use GPU if available)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

Using device: cuda

# Custom transformation to add noise by dropping pixels with a given probability for each channel independently
class AddNoiseTransform:
  def __init__(self, drop_prob=0.1):
    self.drop_prob = drop_prob

  def __call__(self, img):
    channel_mask = (torch.rand(img.shape) > self.drop_prob).float()
    noisy_img = img * channel_mask
    return noisy_img

# PairDataset class that returns a pair of images: (original, noisy)
class PairDataset(Dataset):
  def __init__(self, dataset_origin, dataset_noisy):
    assert len(dataset_origin) == len(dataset_noisy), 'Datasets must be of the same length'
    self.dataset_origin = dataset_origin
    self.dataset_noisy = dataset_noisy

  def __getitem__(self, index):
    clean_img, _ = self.dataset_origin[index]
    noisy_img, _ = self.dataset_noisy[index]
    return clean_img, noisy_img

  def __len__(self):
    return len(self.dataset_origin)

# Function to create paired dataset for specific classes
def create_paired_dataloader(classes, num_samples=1500, batch_size=64, drop_prob=0.5, split='train'):
  """
  Creates a DataLoader for a paired dataset with noisy and clean images from specified classes.

  Parameters:
  - classes (list): List of class labels to include (e.g., [0, 1, 2, 3, 4] or [5, 6, 7, 8, 9]).
  - num_samples (int): Number of samples to use from the filtered classes.
  - batch_size (int): Batch size for the DataLoader.
  - drop_prob (float): Drop probability for AddNoiseTransform.
  - split (str): Dataset split to load ('train' or 'test').

  Returns:
  - DataLoader: A DataLoader for the paired dataset.
  """

  # Define transformations
  transform_clean = transforms.ToTensor()
  transform_noisy = transforms.Compose([transforms.ToTensor(), AddNoiseTransform(drop_prob=drop_prob)])

  # Load the specified dataset split
  dataset_clean_full = datasets.SVHN(root='./data', split=split, download=True, transform=transform_clean)
  dataset_noisy_full = datasets.SVHN(root='./data', split=split, download=True, transform=transform_noisy)

  # Filter to include only samples with specified classes
  indices = [i for i, (_, label) in enumerate(dataset_clean_full) if label in classes]

  # Take only the first `num_samples` samples after filtering
  subset_indices = indices[:num_samples]

  # Create clean and noisy subsets
  dataset_clean = Subset(dataset_clean_full, subset_indices)
  dataset_noisy = Subset(dataset_noisy_full, subset_indices)

  # Create the paired dataset
  paired_dataset = PairDataset(dataset_clean, dataset_noisy)

  # DataLoader for paired dataset
  paired_dataloader = DataLoader(paired_dataset, batch_size=batch_size, shuffle=True)

  return paired_dataloader

class ImageDenoisingCNN(nn.Module, PretrainedMixin):
  def __init__(self):
    super(ImageDenoisingCNN, self).__init__()

    # First Convolutional Layer
    # Input: 32x32x3 -> Output: 32x32x30
    self.conv1 = nn.Conv2d(in_channels=3, out_channels=30, kernel_size=3, padding=1, stride=1)
    self.relu = nn.ReLU()

    # Second Convolutional Layer
    # Input: 32x32x30 -> Output: 32x32x3
    self.conv2 = nn.Conv2d(in_channels=30, out_channels=3, kernel_size=3, padding=1, stride=1)
    self.sigmoid = nn.Sigmoid()

  def forward(self, x):
    # First conv layer with ReLU
    x = self.conv1(x)
    x = self.relu(x)

    # Second conv layer with Sigmoid
    x = self.conv2(x)
    x = self.sigmoid(x)

    return x

model = ImageDenoisingCNN()

# Calculate and print the number of trainable parameters
num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print('Total number of trainable parameters in ImageDenoisingCNN:', num_params)

Total number of trainable parameters in ImageDenoisingCNN: 1653

# Display a single row of images with specified columns
def show_images_grid2(images, title, cols):
  fig, axes = plt.subplots(1, cols, figsize=(cols * 2, 2))
  if cols == 1:
    axes = [axes]

  for idx in range(cols):
    if idx < len(images):
      img = images[idx].numpy().transpose((1, 2, 0))
      axes[idx].imshow(img)
      axes[idx].set_title(title)
      axes[idx].axis('off')

  plt.tight_layout()
  plt.show()

def train(train_loader, test_loader, model, epochs, loss_function, optimizer, device='cuda'):
  """
  Train the model on the training dataset and evaluate it on the test dataset.
  """
  # Move model to the specified device
  model = model.to(device)
  train_loss_epochs = []
  test_loss_epochs = []

  for epoch in range(epochs):
    model.train()
    train_loss_batches = []

    # Use context manager for batch progress bar
    with tqdm(
      enumerate(train_loader), total=len(train_loader), desc=f'epoch {epoch + 1}/{epochs}', ncols=100
    ) as batch_pbar:
      for batch_idx, (clean_images, noisy_images) in batch_pbar:
        # Move data to device
        clean_images = clean_images.to(device)
        noisy_images = noisy_images.to(device)

        # Zero the gradients
        optimizer.zero_grad()

        # Forward pass
        denoised_images = model(noisy_images)
        loss = loss_function(denoised_images, clean_images)

        # Backward pass and optimize
        loss.backward()
        optimizer.step()

        # Track batch loss
        train_loss_batches.append(loss.item())
        batch_pbar.set_postfix({'batch_loss': loss.item()})

        # Display sample results every 5 epochs, at the last batch
        if epoch % 5 == 0 and batch_idx == len(train_loader) - 1:
          show_images_grid2(clean_images[:5].detach().cpu(), title='Clean', cols=5)
          show_images_grid2(noisy_images[:5].detach().cpu(), title='Noisy', cols=5)
          show_images_grid2(denoised_images[:5].detach().cpu(), title='Denoised', cols=5)

    # Calculate average training loss for the epoch
    train_loss_epoch = np.mean(train_loss_batches)
    train_loss_epochs.append(train_loss_epoch)

    # Evaluate model on test set
    test_loss_epoch = evaluate(test_loader, model, loss_function, epoch + 1, num_epochs, device=device)
    test_loss_epochs.append(test_loss_epoch)

  return train_loss_epochs, test_loss_epochs

def evaluate(dataloader, model, loss_function, epoch, num_epochs, device='cuda'):
  """
  Evaluate the model on the test dataset and return the average loss.
  """
  model.eval()
  test_losses = []

  with torch.no_grad():
    with tqdm(dataloader, desc=f'eval  {epoch}/{num_epochs}', ncols=100) as eval_pbar:
      for clean_images, noisy_images in eval_pbar:
        # Move data to device
        clean_images = clean_images.to(device)
        noisy_images = noisy_images.to(device)

        # Forward pass
        denoised_images = model(noisy_images)
        loss = loss_function(denoised_images, clean_images)

        # Track batch loss
        test_losses.append(loss.item())

  return np.mean(test_losses)

learning_rate = 0.05
batch_size = 64
drop_rate = 0.3
num_samples = 1500
num_epochs = 100

# Setup code
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Initialize model, criterion, and optimizer
model = ImageDenoisingCNN().to(device)
criterion = nn.MSELoss()  # Mean Squared Error loss for image reconstruction
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# Create data loaders for first 5 classes (0-4)
first_five_classes = list(range(5))  # [0, 1, 2, 3, 4]
paired_trainloader_first_five = create_paired_dataloader(
  classes=first_five_classes, num_samples=num_samples, batch_size=batch_size, drop_prob=drop_rate, split='train'
)

paired_testloader_first_five = create_paired_dataloader(
  classes=first_five_classes, num_samples=num_samples, batch_size=batch_size, drop_prob=drop_rate, split='test'
)

# Train the model
print('Starting training...')
train_loss_epochs, test_loss_epochs = train(
  paired_trainloader_first_five, paired_testloader_first_five, model, num_epochs, criterion, optimizer, device
)

Using downloaded and verified file: ./data/train_32x32.mat
Using downloaded and verified file: ./data/train_32x32.mat
Using downloaded and verified file: ./data/test_32x32.mat
Using downloaded and verified file: ./data/test_32x32.mat
Starting training...

epoch 1/100:  88%|███████████████████████████▏   | 21/24 [00:00<00:00, 67.47it/s, batch_loss=0.0138]

epoch 1/100: 100%|███████████████████████████████| 24/24 [00:00<00:00, 34.06it/s, batch_loss=0.0138]
eval  1/100: 100%|██████████████████████████████████████████████████| 24/24 [00:00<00:00, 78.86it/s]
epoch 2/100: 100%|██████████████████████████████| 24/24 [00:00<00:00, 69.86it/s, batch_loss=0.00674]
eval  2/100: 100%|██████████████████████████████████████████████████| 24/24 [00:00<00:00, 78.98it/s]
epoch 3/100: 100%|██████████████████████████████| 24/24 [00:00<00:00, 68.84it/s, batch_loss=0.00603]
eval  3/100: 100%|██████████████████████████████████████████████████| 24/24 [00:00<00:00, 78.88it/s]
epoch 4/100: 100%|██████████████████████████████| 24/24 [00:00<00:00, 69.63it/s, batch_loss=0.00606]
eval  4/100: 100%|██████████████████████████████████████████████████| 24/24 [00:00<00:00, 78.44it/s]
epoch 5/100: 100%|██████████████████████████████| 24/24 [00:00<00:00, 69.30it/s, batch_loss=0.00501]
eval  5/100: 100%|██████████████████████████████████████████████████| 24/24 [00:00<00:00, 79.14it/s]
epoch 6/100:  88%|██████████████████████████▎   | 21/24 [00:00<00:00, 69.40it/s, batch_loss=0.00683]

epoch 6/100: 100%|██████████████████████████████| 24/24 [00:00<00:00, 34.39it/s, batch_loss=0.00683]
eval  6/100: 100%|██████████████████████████████████████████████████| 24/24 [00:00<00:00, 79.16it/s]
epoch 7/100: 100%|██████████████████████████████| 24/24 [00:00<00:00, 67.89it/s, batch_loss=0.00721]
eval  7/100: 100%|██████████████████████████████████████████████████| 24/24 [00:00<00:00, 78.22it/s]
epoch 8/100: 100%|██████████████████████████████| 24/24 [00:00<00:00, 69.38it/s, batch_loss=0.00439]
eval  8/100: 100%|██████████████████████████████████████████████████| 24/24 [00:00<00:00, 78.08it/s]
epoch 9/100: 100%|██████████████████████████████| 24/24 [00:00<00:00, 69.79it/s, batch_loss=0.00523]
eval  9/100: 100%|██████████████████████████████████████████████████| 24/24 [00:00<00:00, 79.28it/s]
epoch 10/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 68.22it/s, batch_loss=0.00595]
eval  10/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 78.78it/s]
epoch 11/100:  88%|█████████████████████████▍   | 21/24 [00:00<00:00, 68.39it/s, batch_loss=0.00472]

epoch 11/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 34.40it/s, batch_loss=0.00472]
eval  11/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 79.16it/s]
epoch 12/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 68.70it/s, batch_loss=0.00628]
eval  12/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 78.32it/s]
epoch 13/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 70.69it/s, batch_loss=0.00538]
eval  13/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 80.19it/s]
epoch 14/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 70.71it/s, batch_loss=0.00437]
eval  14/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 79.43it/s]
epoch 15/100: 100%|██████████████████████████████| 24/24 [00:00<00:00, 70.75it/s, batch_loss=0.0044]
eval  15/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 80.31it/s]
epoch 16/100:  92%|██████████████████████████▌  | 22/24 [00:00<00:00, 69.51it/s, batch_loss=0.00433]

epoch 16/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 34.89it/s, batch_loss=0.00433]
eval  16/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 79.53it/s]
epoch 17/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 70.90it/s, batch_loss=0.00498]
eval  17/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 78.70it/s]
epoch 18/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 70.07it/s, batch_loss=0.00435]
eval  18/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 79.04it/s]
epoch 19/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 70.30it/s, batch_loss=0.00425]
eval  19/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 79.67it/s]
epoch 20/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 70.51it/s, batch_loss=0.00375]
eval  20/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 79.25it/s]
epoch 21/100:  67%|████████████████████          | 16/24 [00:00<00:00, 70.11it/s, batch_loss=0.0033]

epoch 21/100: 100%|██████████████████████████████| 24/24 [00:00<00:00, 34.96it/s, batch_loss=0.0033]
eval  21/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 79.22it/s]
epoch 22/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 70.83it/s, batch_loss=0.00432]
eval  22/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 78.70it/s]
epoch 23/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 68.46it/s, batch_loss=0.00363]
eval  23/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 78.22it/s]
epoch 24/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 70.51it/s, batch_loss=0.00387]
eval  24/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 79.92it/s]
epoch 25/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 71.31it/s, batch_loss=0.00339]
eval  25/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 78.97it/s]
epoch 26/100:  88%|█████████████████████████▍   | 21/24 [00:00<00:00, 68.08it/s, batch_loss=0.00419]

epoch 26/100: 100%|█████████████████████████████| 24/24 [00:01<00:00, 21.25it/s, batch_loss=0.00419]
eval  26/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 72.61it/s]
epoch 27/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 61.17it/s, batch_loss=0.00372]
eval  27/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 71.89it/s]
epoch 28/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 70.79it/s, batch_loss=0.00317]
eval  28/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 77.65it/s]
epoch 29/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 67.36it/s, batch_loss=0.00296]
eval  29/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 77.80it/s]
epoch 30/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 69.64it/s, batch_loss=0.00253]
eval  30/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 78.53it/s]
epoch 31/100:  88%|█████████████████████████▍   | 21/24 [00:00<00:00, 69.38it/s, batch_loss=0.00385]

epoch 31/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 34.47it/s, batch_loss=0.00385]
eval  31/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 79.06it/s]
epoch 32/100: 100%|██████████████████████████████| 24/24 [00:00<00:00, 69.42it/s, batch_loss=0.0049]
eval  32/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 79.53it/s]
epoch 33/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 71.00it/s, batch_loss=0.00414]
eval  33/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 77.90it/s]
epoch 34/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 69.02it/s, batch_loss=0.00431]
eval  34/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 78.77it/s]
epoch 35/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 68.26it/s, batch_loss=0.00469]
eval  35/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 76.89it/s]
epoch 36/100:  88%|█████████████████████████▍   | 21/24 [00:00<00:00, 65.60it/s, batch_loss=0.00452]

epoch 36/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 33.07it/s, batch_loss=0.00452]
eval  36/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 73.71it/s]
epoch 37/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 60.16it/s, batch_loss=0.00346]
eval  37/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 74.22it/s]
epoch 38/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 60.31it/s, batch_loss=0.00322]
eval  38/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 76.56it/s]
epoch 39/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 69.06it/s, batch_loss=0.00401]
eval  39/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 79.02it/s]
epoch 40/100: 100%|█████████████████████████████| 24/24 [00:00<00:00, 70.78it/s, batch_loss=0.00301]
eval  40/100: 100%|█████████████████████████████████████████████████| 24/24 [00:00<00:00, 78.72it/s]
epoch 41/100:  88%|█████████████████████████▍   | 21/24 [00:00<00:00, 68.81it/s, batch_loss=0.00295]

# Create the plot
plt.figure(figsize=(10, 6))

# Plot training and test losses
epochs = range(1, len(train_loss_epochs) + 1)
plt.plot(epochs, train_loss_epochs, label='Training Loss', color='blue', linestyle='-')
plt.plot(epochs, test_loss_epochs, label='Test Loss', color='red', linestyle='-')

# Customize the plot
plt.title('Training and Test Losses Over Time', fontsize=14, pad=15)
plt.xlabel('Epochs', fontsize=12)
plt.ylabel('Loss (MSE)', fontsize=12)
plt.grid(True, linestyle='--', alpha=0.7)
plt.legend(fontsize=10)

# Add minor gridlines
plt.minorticks_on()
plt.grid(True, which='minor', linestyle=':', alpha=0.4)

# Adjust layout and display
plt.tight_layout()
plt.show()

# Print final losses
print(f'Final Training Loss: {train_loss_epochs[-1]:.6f}')
print(f'Final Test Loss: {test_loss_epochs[-1]:.6f}')

Final Training Loss: 0.003326
Final Test Loss: 0.003811

# Load test data for classes 5-9 with 1,500 samples per class
last_five_classes = list(range(5, 10))  # [5, 6, 7, 8, 9]
paired_testloader_last_five = create_paired_dataloader(
  classes=last_five_classes, num_samples=num_samples, batch_size=batch_size, drop_prob=drop_rate, split='test'
)

print(f'Number of batches in paired_testloader_last_five: {len(paired_testloader_last_five)}')

# Evaluate model on the last 5 classes test set
model.eval()
test_loss = 0
total_batches = 0

# Store some examples for visualization
example_clean = None
example_noisy = None
example_denoised = None

with torch.no_grad():
  for batch_idx, (clean_images, noisy_images) in enumerate(paired_testloader_last_five):
    # Move images to device
    clean_images = clean_images.to(device)
    noisy_images = noisy_images.to(device)

    # Get model predictions
    denoised_images = model(noisy_images)

    # Calculate loss
    batch_loss = criterion(denoised_images, clean_images)
    test_loss += batch_loss.item()
    total_batches += 1

    # Store the first batch for visualization
    if batch_idx == 0:
      example_clean = clean_images
      example_noisy = noisy_images
      example_denoised = denoised_images

# Calculate average test loss
average_test_loss = test_loss / total_batches
print(f'Average Test Loss on classes 5-9: {average_test_loss:.6f}')

Using downloaded and verified file: ./data/test_32x32.mat
Using downloaded and verified file: ./data/test_32x32.mat
Number of batches in paired_testloader_last_five: 24
Average Test Loss on classes 5-9: 0.003754

# Move images to CPU for visualization and detach from computational graph
example_clean = example_clean.cpu().detach()
example_noisy = example_noisy.cpu().detach()
example_denoised = example_denoised.cpu().detach()


show_images_grid2(example_clean[:5], title='Clean Images', cols=5)
show_images_grid2(example_noisy[:5], title='Noisy Images', cols=5)
show_images_grid2(example_denoised[:5], title='Denoised Images', cols=5)

plt.tight_layout()
plt.show()

# Save example images if needed
save_path = 'denoising_examples_last_five.png'
plt.savefig(save_path)
print(f'Example images saved to {save_path}')

<Figure size 640x480 with 0 Axes>

Example images saved to denoising_examples_last_five.png

<Figure size 640x480 with 0 Axes>

Assignment 4: Application of Neural Networks for Image Processing¶

SVHN Dataset¶

Visualization of SVHN dataset¶

Task 1: SVHN Image Classification Using a Convolutional Neural Network (CNN) (30 points) ¶

A (15 points) Build the CNN Model¶

B (15 points) Train the network¶

Task 2: SVHN Denoising with Convolution Neural Network (CNN) (45 points) ¶

Create the Noisy dataset¶

Step 1: Defining a Custom Noise Transformation¶

Step 2: Creating a Dataset of Paired Images: Original and Noisy¶

Task A (10 points): Implementing a CNN for Image Denoising¶

Architecture Details¶

Implementation Requirements¶

Task B (20 points): Denoising SVHN Images¶

Task C (5 points): Plotting Training and Testing Losses Over Epochs¶

Task D (10 points): Denoising Last 5 Classes, Reporting Loss, and Visualization¶

Task 3: Kaggle Competition (25 + 20 Bonus Points) ¶