How to Implement Custom Loss Functions or Metrics in YOLOv5

You can implement custom loss functions or metrics in YOLOv5 by subclassing nn.Module for losses in utils/loss.py or adding metric functions to utils/metrics.py, then integrating them into the ComputeLoss class or validation loop without modifying the core training pipeline.

YOLOv5's modular architecture separates loss computation and metric evaluation into dedicated utility modules, making it straightforward to implement custom loss functions or metrics in YOLOv5. The repository organizes training logic in utils/loss.py through the ComputeLoss class, while validation metrics reside in utils/metrics.py. This design allows you to inject custom PyTorch modules or evaluation functions while preserving the existing training loop in train.py.

Understanding YOLOv5's Loss and Metric Architecture

Loss Computation in utils/loss.py

The loss computation is encapsulated in utils/loss.py by the ComputeLoss class. It gathers the model's predictions, builds the target tensors, and combines several sub-losses for classification, objectness, and box regression. Each sub-loss is a regular torch.nn.Module, such as BCEBlurWithLogitsLoss, FocalLoss, or QFocalLoss.

Metric Evaluation in utils/metrics.py

Metric evaluation is performed in utils/metrics.py during validation. Functions such as fitness, ap_per_class, and ConfusionMatrix compute precision/recall, mAP, and class-wise scores. These functions accept model predictions and ground-truth tensors, making them easy to extend.

Implementing a Custom Loss Function in YOLOv5

Step 1: Create the Loss Module

Create a new loss class by subclassing nn.Module and implementing the forward method. Add this to utils/loss.py or import it from a new file.

# utils/loss.py – add after the existing loss classes

class DiceLoss(nn.Module):
    """
    Dice loss for segmentation masks.
    Reduces the overlap error between predicted mask (logits) and target mask.
    """
    def __init__(self, eps: float = 1e-6):
        super().__init__()
        self.eps = eps

    def forward(self, pred, true):
        """
        pred: logits (N, C, H, W) – raw model output
        true: binary mask (N, C, H, W) – same shape as pred
        """
        prob = torch.sigmoid(pred)
        num = 2.0 * (prob * true).sum(dim=(2, 3))
        den = prob.sum(dim=(2, 3)) + true.sum(dim=(2, 3)) + self.eps
        loss = 1.0 - num / den
        return loss.mean()

Step 2: Integrate into ComputeLoss

Instantiate your loss class inside ComputeLoss.__init__ and call it within the __call__ method to aggregate it into the total loss.

# utils/loss.py – modify ComputeLoss.__init__

self.dice = DiceLoss()  # ← new loss instance

# utils/loss.py – inside __call__

mask_logits = p[-1]  # example: last head contains mask logits

mask_targets = targets[..., 5:]  # adapt to your label layout

lseg = self.dice(mask_logits, mask_targets)
lseg = lseg * self.hyp.get("seg", 1.0)  # optional weighting via hyper-params

# aggregate

lbox += lseg

Step 3: Configure Hyperparameters

If your loss requires tunable weights, extend a hyper-parameter YAML file (e.g., data/hyps/hyp.scratch-low.yaml) and read the value via the hyp dictionary.

# data/hyps/hyp.scratch-low.yaml

seg: 1.0  # weight for the segmentation Dice loss

train.py loads opt.hyp and passes it to ComputeLoss, making the value accessible as self.hyp["seg"].

Implementing a Custom Metric in YOLOv5

Step 1: Define the Metric Function

Add your metric function to utils/metrics.py. It should accept predictions and ground-truth tensors with the same signature as existing functions like ap_per_class.

# utils/metrics.py – add at the end of the file

def mean_iou(pred_masks, true_masks, threshold: float = 0.5):
    """
    Computes mean Intersection-over-Union across a batch.
    pred_masks – raw logits (N, C, H, W)
    true_masks – binary ground truth (N, C, H, W)
    """
    pred = (torch.sigmoid(pred_masks) > threshold).float()
    inter = (pred * true_masks).sum(dim=(2, 3))
    union = pred.sum(dim=(2, 3)) + true_masks.sum(dim=(2, 3)) - inter
    iou = inter / (union + 1e-7)
    return iou.mean().item()

Step 2: Integrate into Validation Loop

Call your metric function from val.py (or segment/val.py for segmentation models) after the model produces predictions, then log the result via the existing logger.

# segment/val.py – after `pred = model(imgs)`

mask_pred = pred[-1]  # adjust index for your model

mask_iou = mean_iou(mask_pred, targets[..., 5:])
LOGGER.info(f"Mean IoU (mask): {mask_iou:.4f}")

Logging Custom Metrics to Experiment Trackers

To ensure your custom metric appears in TensorBoard, Weights & Biases, or CSV logs, append it to the log_vals list in train.py around line 78–79:

# train.py – after computing `results, maps, _ = validate.run(...)`

log_vals.extend([mask_iou])  # append custom metric

This writes the metric through all active loggers without requiring additional instrumentation.

Summary

• Create custom losses by subclassing nn.Module in utils/loss.py and integrating them into the ComputeLoss class.
• Add custom metrics as functions in utils/metrics.py and invoke them from val.py or segment/val.py.
• Configure weights via YAML hyper-parameter files (e.g., data/hyps/hyp.scratch-low.yaml) and access them through the hyp dictionary.
• Log metrics by appending values to log_vals in train.py to ensure visibility in TensorBoard, W&B, and CSV outputs.

Frequently Asked Questions

Where do I place custom loss classes in YOLOv5?

Place custom loss classes in utils/loss.py alongside existing losses like FocalLoss and BCEBlurWithLogitsLoss, or create a new file and import the class into utils/loss.py. The ComputeLoss class must be able to instantiate your module during its __init__ method.

Can I use custom losses for segmentation tasks in YOLOv5?

Yes. For segmentation models using segment/train.py, the ComputeLoss class receives mask predictions in the final elements of the prediction list. You can apply custom segmentation losses like DiceLoss or TverskyLoss to these mask logits by extracting p[-1] for predictions and the corresponding target mask channels from the labels tensor.

How do I weight multiple custom losses in YOLOv5?

Add weighting coefficients to a hyper-parameter YAML file in data/hyps/ (e.g., seg: 0.5 or custom_loss_alpha: 1.0). The train.py script loads these into the hyp dictionary and passes it to ComputeLoss. Inside ComputeLoss.__call__, retrieve the weight with self.hyp.get("key", default_value) and multiply it by the loss tensor before adding it to the total loss.

Will custom metrics affect model training in YOLOv5?

No. Custom metrics added to utils/metrics.py and called from val.py operate during the validation phase only. They compute statistics on predictions versus ground truth to monitor model performance but do not contribute gradients to the optimization loop. Only losses integrated into ComputeLoss and called during the training forward pass affect weight updates.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}