Training.py

import torch
import torch.nn as nn
from torchvision import transforms
import sys
import os
from pycocotools.coco import COCO
from data_loader import get_loader
from model import EncoderCNN, DecoderRNN
import math
import numpy as np

batch_size = 128          # batch size
vocab_threshold = 5        # minimum word count threshold
vocab_from_file = False    # if True, load existing vocab file
embed_size = 256           # dimensionality of image and word embeddings
hidden_size = 512          # number of features in hidden state of the RNN decoder
num_epochs = 3             # number of training epochs
save_every = 1             # determines frequency of saving model weights
print_every = 100          # determines window for printing average loss
log_file = r'C:\Users\Shabnam\Desktop\FrameToPhrase\logs\training_log.txt'

transform_train = transforms.Compose([ 
    transforms.Resize(256),                          # smaller edge of image resized to 256
    transforms.RandomCrop(224),                      # get 224x224 crop from random location
    transforms.RandomHorizontalFlip(),               # horizontally flip image with probability=0.5
    transforms.ToTensor(),                           # convert the PIL Image to a tensor
    transforms.Normalize((0.485, 0.456, 0.406),      # normalize image for pre-trained model
                         (0.229, 0.224, 0.225))])

data_loader = get_loader(transform=transform_train,
                         mode='train',
                         batch_size=batch_size,
                         vocab_threshold=vocab_threshold,
                         vocab_from_file=vocab_from_file)

# The size of the vocabulary.
vocab_size = len(data_loader.dataset.vocab)

# Initialize the encoder and decoder. 
encoder = EncoderCNN(embed_size)
decoder = DecoderRNN(embed_size, hidden_size, vocab_size)

# Move models to GPU if CUDA is available. 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
encoder.to(device)
decoder.to(device)

# Define the loss function. 
criterion = nn.CrossEntropyLoss().cuda() if torch.cuda.is_available() else nn.CrossEntropyLoss()

# Specify the learnable parameters of the model.
params = list(decoder.parameters()) + list(encoder.embed.parameters())

# Define the optimizer.
optimizer = torch.optim.Adam(params, lr=0.001)

# Set the total number of training steps per epoch.
total_step = math.ceil(len(data_loader.dataset.caption_lengths) / data_loader.batch_sampler.batch_size)

# Create directories if they don't exist
os.makedirs('./models', exist_ok=True)
os.makedirs(os.path.dirname(log_file), exist_ok=True)

# Open the training log file.
f = open(log_file, 'w')

print(f"Starting training for {num_epochs} epochs...")
print(f"Total steps per epoch: {total_step}")
print(f"Device: {device}")

# TRAINING LOOP - THIS WAS MISSING!
for epoch in range(1, num_epochs + 1):
    
    total_loss = 0.0
    
    for i_step in range(1, total_step + 1):
        
        # Randomly sample a caption length, and sample indices with that length.
        indices = data_loader.dataset.get_train_indices()
        # Create and assign a batch sampler to retrieve a batch with the sampled indices.
        new_sampler = torch.utils.data.sampler.SubsetRandomSampler(indices=indices)
        data_loader.batch_sampler.sampler = new_sampler
        
        # Obtain the batch.
        images, captions = next(iter(data_loader))
        
        # Move batch to device
        images = images.to(device)
        captions = captions.to(device)
        
        # Zero the gradients.
        decoder.zero_grad()
        encoder.zero_grad()
        
        # Pass the inputs through the CNN-RNN model.
        features = encoder(images)
        outputs = decoder(features, captions)
        
        # Calculate the batch loss.
        loss = criterion(outputs.view(-1, vocab_size), captions.view(-1))
        
        # Backward pass.
        loss.backward()
        
        # Update the parameters in the optimizer.
        optimizer.step()
        
        # Get training statistics.
        total_loss += loss.item()
        stats = f'Epoch [{epoch}/{num_epochs}], Step [{i_step}/{total_step}], Loss: {loss.item():.4f}, Perplexity: {np.exp(loss.item()):.4f}'
        
        # Print training statistics (on same line).
        if i_step % print_every == 0:
            print('\r' + stats)
            
        # Log training statistics to file.
        f.write(stats + '\n')
        f.flush()
        
    # Save the weights after each epoch
    if epoch % save_every == 0:
        torch.save(decoder.state_dict(), os.path.join('./models', f'decoder-{epoch}.pkl'))
        torch.save(encoder.state_dict(), os.path.join('./models', f'encoder-{epoch}.pkl'))
        print(f'\nSaved models for epoch {epoch}')

# Close the training log file.
f.close()

print("\nTraining complete!")
print("Model files saved in ./models/")