[Python] Object Detection Custom Model Loss Not Decreasing During Training

I'm working on an object detection project using PyTorch on Kaggel, where I'm trying to detect and recognize text in any day objects, such as traffic signs, documents, etc. I've implemented a custom model based on the EAST model architecture, using ResNet as my backbone. However, I'm encountering an issue where the loss isn't decreasing significantly during training. Here's a snippet of my training output:

...
Epoch 28/30
Training: |####################| 005/005 [100.00%] in 0.0s (1746419.0/s, ETA: 0.0s) (loss: 53.674442291259766)
Validation: |####################| 003/003 [100.00%] in 0.0s (731885.0/s, ETA: 0.0s) (loss: 52.65853627522787)
Epoch 29/30
Training: |####################| 005/005 [100.00%] in 0.0s (1104972.0/s, ETA: 0.0s) (loss: 54.1553955078125)
Validation: |####################| 003/003 [100.00%] in 0.0s (950871.0/s, ETA: 0.0s) (loss: 52.0231081644694)
Epoch 30/30
Training: |####################| 005/005 [100.00%] in 0.0s (1353912.0/s, ETA: 0.0s) (loss: 53.528507232666016)
Validation: |####################| 003/003 [100.00%] in 0.0s (1148105.0/s, ETA: 0.0s) (loss: 51.57835261027018)


I've tried the following without much success:

• Adjusting the learning rate
• Changing the batch size
• Modifying the model architecture slightly

It doesn't make sense to me, I have assume I did everything correctly in my module, I've split apart my two detections, and return a tensor, in which, I can use it on my custom loss class and calculate it.

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torchvision.models as models
import torchvision.transforms as transforms
from typing import Tuple, Optional, List
import logging
from pathlib import Path

# Set up logging configuration
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(levelname)s - %(message)s",
handlers=[
logging.FileHandler("model.log", mode="w")
]
)

logger = logging.getLogger(__name__)

class TLR(nn.Module):
"""
Text Localization and Recognition (TLR) model with additional layers.

This model uses a ResNet50 backbone and adds more convolutional layers
for improved feature extraction and prediction.
"""

def __init__(self, hidden_size: int = 256, input_channels: int = 1, num_classes: int = 35):
"""
Initialize the TLR model.

Args:
hidden_size (int): Size of the hidden layers. Defaults to 256.
input_channels (int): Number of input channels. Defaults to 1.
num_classes (int): Number of classes to predict. Defaults to 35.
"""
super(TLR, self).__init__()

logger.info("Initializing TLR model")

self.num_classes = num_classes

# Initialize ResNet50 backbone without pretrained weights
backbone = models.resnet50(weights=None)
backbone.conv1 = nn.Conv2d(input_channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)

self.backbone = nn.Sequential(*list(backbone.children())[:-2])

self.merge = nn.Sequential(
nn.Conv2d(2048, hidden_size, 1),
nn.BatchNorm2d(hidden_size),
nn.ReLU(inplace=True),
nn.Conv2d(hidden_size, hidden_size, 3, padding=1),
nn.BatchNorm2d(hidden_size),
nn.ReLU(inplace=True),
nn.Conv2d(hidden_size, hidden_size, 3, padding=1),
nn.BatchNorm2d(hidden_size),
nn.ReLU(inplace=True)
)

self.det_head = nn.Sequential(
nn.Conv2d(hidden_size, hidden_size, 3, padding=1),
nn.BatchNorm2d(hidden_size),
nn.ReLU(inplace=True),
nn.Conv2d(hidden_size, hidden_size // 2, 1),
nn.BatchNorm2d(hidden_size // 2),
nn.ReLU(inplace=True),
nn.Conv2d(hidden_size // 2, hidden_size // 2, 3, padding=1),
nn.BatchNorm2d(hidden_size // 2),
nn.ReLU(inplace=True),
nn.Conv2d(hidden_size // 2, 4 + num_classes, 1)
)

logger.info("EnhancedTLR model initialized")

def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Forward pass of the EnhancedTLR model.

Args:
x (torch.Tensor): Input tensor of shape [batch_size, input_channels, height, width]

Returns:
Tuple[torch.Tensor, torch.Tensor]:
- bbox: Bounding box predictions of shape [batch_size, num_predictions, 4]
- class_scores: Class scores of shape [batch_size, num_predictions, num_classes]
"""
features = self.backbone(x)
logger.debug(f"Backbone output shape: {features.shape}")

merged = self.merge(features)
logger.debug(f"Merge output shape: {merged.shape}")

det_output = self.det_head(merged)
logger.debug(f"Detection head output shape: {det_output.shape}")

batch_size, _, height, width = det_output.shape
det_output = det_output.permute(0, 2, 3, 1).contiguous()
det_output = det_output.view(batch_size, -1, 4 + self.num_classes)

bbox = torch.sigmoid(det_output[..., :4])
class_scores = det_output[..., 4:]

logger.debug(f"Output shapes - bbox: {bbox.shape}, class_scores: {class_scores.shape}")

return bbox, class_scores

class TLRLoss(nn.Module):
"""
Loss function for the TLR model, combining Smooth L1 Loss for bounding box regression
and Cross Entropy Loss for class prediction.
"""
def __init__(self, beta: float = 1.0):
"""
Initialize the TLRLoss.
Args:
beta (float): Transition point from L1 to L2 loss. Defaults to 1.0.
"""
super(TLRLoss, self).__init__()
self.bbox_loss_fn = nn.SmoothL1Loss(beta=beta, reduction='sum')
self.class_loss_fn = nn.CrossEntropyLoss(ignore_index=-1, reduction='sum') # Ignore padding with -1

def forward(
self,
bbox_pred: torch.Tensor,
bbox_target: List[torch.Tensor],
class_scores_pred: torch.Tensor,
class_target: List[torch.Tensor]
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Compute the bounding box regression loss and class classification loss.

Args:
bbox_pred (torch.Tensor): Predicted bounding boxes of shape [batch_size, num_predictions, 4].
bbox_target (List[torch.Tensor]): Ground truth bounding boxes for each image,
where each tensor has shape [num_boxes, 4].
class_scores_pred (torch.Tensor): Predicted class scores of shape [batch_size, num_predictions, num_classes].
class_target (List[torch.Tensor]): Ground truth class labels for each image,
where each tensor has shape [num_boxes].

Returns:
Tuple[torch.Tensor, torch.Tensor]: The bounding box regression loss and class classification loss.
"""
batch_size = bbox_pred.size(0)
device = bbox_pred.device

# Compute total number of boxes across the batch
total_boxes = sum(boxes.shape[0] for boxes in bbox_target)

# Initialize tensors to store all ground truth boxes and labels
all_bbox_target = torch.zeros(total_boxes, 4, device=device)
all_class_target = torch.full((total_boxes,), -1, dtype=torch.long, device=device)

# Initialize mask for valid predictions
valid_pred_mask = torch.zeros(batch_size, bbox_pred.size(1), dtype=torch.bool, device=device)

# Fill in the ground truth tensors and create the mask
start_idx = 0
for i, (boxes, labels) in enumerate(zip(bbox_target, class_target)):
num_boxes = boxes.shape[0]
all_bbox_target[start_idx:start_idx + num_boxes] = boxes
all_class_target[start_idx:start_idx + num_boxes] = labels
valid_pred_mask[i, :num_boxes] = True
start_idx += num_boxes

# Flatten predictions and apply mask
bbox_pred_flat = bbox_pred[valid_pred_mask]
class_scores_pred_flat = class_scores_pred[valid_pred_mask]

# Compute losses
bbox_loss = self.bbox_loss_fn(bbox_pred_flat, all_bbox_target)
class_loss = self.class_loss_fn(class_scores_pred_flat, all_class_target)

# Normalize losses by the total number of boxes
bbox_loss /= total_boxes
class_loss /= total_boxes

return bbox_loss, class_loss

if __name__ == "__main__":
# Define paths
images_path: Path = Path("/kaggle/input/license-plate-individual-character-recognition/dataset/test/images")
annotations_path: Path = Path("/kaggle/input/license-plate-individual-character-recognition/dataset/test/annotations")

# Define transformations for images
transform: transforms.Compose = transforms.Compose([
transforms.Resize((150, 300)),
transforms.ToTensor()
])

# Initialize dataset and DataLoader
dataset: PairedDataset = PairedDataset(images_path, annotations_path, transform=transform)
dataloader: DataLoader = DataLoader(dataset, batch_size=32, collate_fn=collate_fn)

# Initialize model with custom arguments
model: TLR = TLR(hidden_size=256, input_channels=1, num_classes=35)

# Test the model
model.eval()
with torch.no_grad():
for batch_idx, (images, targets) in enumerate(dataloader):
if images is None:
continue

logger.info(f"Processing batch {batch_idx + 1}/{len(dataloader)}")

bbox_pred, class_pred = model(images)

print(f"Batch {batch_idx + 1}:")
print(f"Bounding box predictions shape: {bbox_pred}")
print(f"Class predictions shape: {class_pred}")

logger.info("Data pass completed")
logger.info("Script execution completed")


In my training script, I split my targets, that are in a list, into their own group, each image has multiple detections. And each bounding box has its own label.

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
import torchvision.models.detection as detection
from pathlib import Path
import numpy as np
import random
import logging
import asyncio
from typing import Union, List, Tuple, Optional, Dict, Literal

# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(levelname)s - %(message)s",
handlers=[
logging.FileHandler("trainer.log", mode="w")
]
)

class Trainer:
"""
A class to handle training, validation, and testing of an object detection model.
"""

def __init__(self, module: nn.Module, train_data_paths: List[Path], val_data_paths: List[Path],
test_data_paths: List[Path], image_size: Tuple[int, int] = (150, 300),
epochs: int = 30, batch_size: int = 64, learning_rate: float = 0.0001,
seed: Optional[int] = None) -> None:
"""
Initialize the Trainer with dataset paths, hyperparameters, and device settings.

Args:
module (nn.Module): The model to use for training.
train_data_paths (List[Path]): Paths to training images and annotations.
val_data_paths (List[Path]): Paths to validation images and annotations.
test_data_paths (List[Path]): Paths to test images and annotations.
image_size (Tuple[int, int]): Desired image size (height, width).
epochs (int): Number of training epochs.
batch_size (int): Batch size for DataLoader.
learning_rate (float): Learning rate for the optimizer.
seed (Optional[int]): Random seed for reproducibility.
"""
# Validate input paths
for paths in [train_data_paths, val_data_paths, test_data_paths]:
if not isinstance(paths, list) or len(paths) != 2 or not all(isinstance(p, Path) for p in paths):
raise TypeError("Each data path argument must be a list of two Path instances")

# Validate hyperparameters
if not isinstance(image_size, tuple) or len(image_size) != 2 or not all(isinstance(i, int) for i in image_size):
raise ValueError("image_size must be a tuple of two integers (height, width)")
if not isinstance(epochs, int) or not isinstance(batch_size, int) or not isinstance(learning_rate, float):
raise TypeError("epochs must be an integer, batch_size must be an integer, and learning_rate must be a float")
if seed is not None and not isinstance(seed, int):
raise TypeError("seed must be an integer or None")
if not isinstance(module, nn.Module):
raise TypeError("module must be an instance of nn.Module")

# Set random seed if provided
self._set_random_seed(seed)

# Set the device to CUDA if available, otherwise use CPU
self.device: torch.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load the model and move to device
self.module: nn.Module = module
self.module.to(self.device)

# Store hyperparameters and prepare DataLoaders
self.image_size: Tuple[int, int] = image_size
self.epochs: int = epochs
self.train_loader: DataLoader = self._load_data(train_data_paths, batch_size, mode="training")
self.val_loader: DataLoader = self._load_data(val_data_paths, batch_size, mode="validation")
self.test_loader: DataLoader = self._load_data(test_data_paths, batch_size, mode="validation")

# Use TLRLoss as the criterion
self.criterion = TLRLoss()

# Setup optimizer and learning rate scheduler
self.optimizer: torch.optim = torch.optim.Adam(self.module.parameters(), lr=learning_rate)
self.scheduler: torch.optim.lr_scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=3, gamma=0.1)

logging.info("Trainer initialized successfully.")

@staticmethod
def _set_random_seed(seed: Optional[int]) -> None:
"""
Set the random seed for reproducibility.

Args:
seed (Optional[int]): Random seed to set. If None, no seed is set.
"""
if seed is not None:
torch.manual_seed(seed)
random.seed(seed)
np.random.seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
logging.info(f"Seed set to {seed}.")

def _load_data(self, data_paths: List[Path], batch_size: int, mode: Literal["training", "validation"] = "training") -> DataLoader:
"""
Load data using the specified paths and return a DataLoader.

Args:
data_paths (List[Path]): Paths to the images and annotations.
batch_size (int): Number of samples per batch.
mode (Literal["training", "validation"]): Mode specifying whether the data is for training or validation.

Returns:
DataLoader: A DataLoader object for the dataset.
"""
transform: transforms.Compose = transforms.Compose([
transforms.Resize(self.image_size),
transforms.RandomRotation(10) if mode == "training" else transforms.Lambda(lambda x: x),
transforms.RandomAffine(degrees=0, translate=(0.1, 0.1), scale=(0.9, 1.1)) if mode == "training" else transforms.Lambda(lambda x: x),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5], std=[0.5])
])
# Initialize the dataset with the given paths and transformations
dataset: PairedDataset = PairedDataset(
images_dir=data_paths[0], annotations_dir=data_paths[1], transform=transform
)
logging.info(f"Data loaded from {data_paths}.")
# Return a DataLoader for the dataset
return DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)

async def _run_epoch_async(self, data_loader: DataLoader, mode: Literal["training", "validation"] = "training") -> float:
"""
Run a single epoch asynchronously for training or validation.

Args:
data_loader (DataLoader): The DataLoader to iterate over.
mode (Literal["training", "validation"]): Mode specifying whether the data is for training or validation.

Returns:
float: The average loss over the epoch.
"""
# Set the model to training or evaluation mode based on the provided mode
self.module.train() if mode == "training" else self.module.eval()

# Initialize the total loss for the epoch
total_loss: float = 0.0

# Initialize a progress bar with asynchronous capabilities
async with Bar(iterations=len(data_loader), title=mode.capitalize(), steps=20) as bar:
start_time: float = asyncio.get_event_loop().time()

# Iterate over the DataLoader, which provides batches of images and targets
for progress, (images, targets) in enumerate(data_loader, start=1):
# Ensure that the images are in the correct format (torch.Tensor)
if not isinstance(images, torch.Tensor):
logging.warning("Images must be a torch tensor.")
continue

# Move images and targets to the device (e.g., GPU)
images = images.to(self.device)
targets: List[Dict[str, torch.Tensor]] = [{k: v.to(self.device) for k, v in t.items()} for t in
targets]

# Reset the gradients of the optimizer
self.optimizer.zero_grad()

# Enable gradient calculation only if in training mode
with torch.set_grad_enabled(mode == "training"):
# Forward pass through the model
bbox_pred, class_pred = self.module(images)

# Scale bounding boxes to original image size (avoid in-place operations)
bbox_pred = torch.stack([
bbox_pred[:, :, 0] * self.image_size[1], # x_min scaled by width
bbox_pred[:, :, 1] * self.image_size[0], # y_min scaled by height
bbox_pred[:, :, 2] * self.image_size[1], # x_max scaled by width
bbox_pred[:, :, 3] * self.image_size[0], # y_max scaled by height
], dim=2)

# Separate targets into bbox_targets and class_targets
bbox_target = [t['boxes'] for t in targets]
class_target = [t['labels'] for t in targets]

# Compute losses for bounding boxes and classes
bbox_loss, class_loss = self.criterion(
bbox_pred,
bbox_target,
class_pred,
class_target
)

loss = (bbox_loss + class_loss)

if mode == "training":
# Backward pass
loss.backward()
self.optimizer.step()

# Accumulate the scalar loss
total_loss += loss.item()

# Await progress updates for the progress bar
await bar.update(progress, start_time)
await bar.postfix(loss=np.divide(total_loss, progress))

# Compute the average loss for the epoch
avg_loss: float = np.divide(total_loss, len(data_loader))
logging.info(f"{mode.capitalize()} epoch completed with average loss: {avg_loss:.4f}")

return avg_loss

async def train(self) -> None:
"""
Train the model for the specified number of epochs.
"""
for epoch in range(self.epochs):
print(f"Epoch {epoch + 1}/{self.epochs}")

logging.info(f"Starting training epoch {epoch + 1}/{self.epochs}.")

# Run the training and validation epochs asynchronously
train_loss: float = await self._run_epoch_async(self.train_loader, mode="training")
logging.info(f"Training epoch {epoch + 1} completed with average loss: {train_loss:.4f}")

val_loss: float = await self._run_epoch_async(self.val_loader, mode="validation")
logging.info(f"Validation epoch {epoch + 1} completed with average loss: {val_loss:.4f}")

# Update the learning rate scheduler
self.scheduler.step()

async def test(self) -> None:
"""
Test the model on the test dataset.
"""
logging.info("Starting testing.")
# Implement testing logic here...

async def save(self, save_path: Union[str, Path]) -> None:
"""
Save the model to the specified path.

Args:
save_path (Union[str, Path]): The path where the model will be saved.
"""
torch.save(self.module, save_path)
print(f"Model saved at: {save_path}")
logging.info(f"Model saved at: {save_path}")

if __name__ == "__main__":
# Create a Trainer instance with dataset paths
trainer = Trainer(
module=TLR(),
train_data_paths=[Path("/kaggle/input/license-plate-individual-character-recognition/dataset/train/images"), Path("/kaggle/input/license-plate-individual-character-recognition/dataset/train/annotations")],
val_data_paths=[Path("/kaggle/input/license-plate-individual-character-recognition/dataset/val/images"), Path("/kaggle/input/license-plate-individual-character-recognition/dataset/val/annotations")],
test_data_paths=[Path("/kaggle/input/license-plate-individual-character-recognition/dataset/test/images"), Path("/kaggle/input/license-plate-individual-character-recognition/dataset/test/annotations")],
)

# Run the training, validating, testing, and saving processes asynchronously
await trainer.train()
await trainer.test()
await trainer.save("recognition_model.pt")


What could be causing this stagnation in the loss, and what strategies would you recommend to improve the training process? Are there any specific debugging steps or modifications to the training loop that might help identify the root cause?

Continue reading...