Building Custom Object Detection Models with PyTorch

Custom object detection is useful when a generic model cannot recognise the objects, environments, or failure conditions that matter to your product. Examples include identifying crop disease in Indian fields, counting vehicles on mixed-traffic roads, spotting defects on factory lines, and reading safety violations from CCTV. PyTorch gives you control over the data pipeline, model components, losses, and deployment path, but that flexibility also makes engineering discipline essential.

This guide uses torchvision and Faster R-CNN as a dependable baseline. The same workflow applies when you later move to RetinaNet, YOLO-based PyTorch implementations, or transformer detectors such as DETR. If you are still deciding how to structure a broader computer-vision project, compare this workflow with how to build computer vision models on GitHub.

Start with the detection problem, not the model

Define the operational question before selecting an architecture. “Detect vehicles” is incomplete; a production requirement might be “detect buses, cars, trucks, and auto-rickshaws in 1080p roadside video, with fewer than one false alert per camera-hour.” That definition determines annotation policy, confidence thresholds, latency targets, and hardware.

Write down:

Classes: Include only distinctions your users or downstream system need.
Object size: Small objects require higher image resolution and stronger feature pyramids.
Scene conditions: Record variations in light, weather, camera angle, crowding, and occlusion.
Latency and throughput: Separate batch processing, near-real-time video, and frame-by-frame edge inference.
Failure cost: A missed safety defect may matter more than an occasional duplicate detection.

For a first implementation, Faster R-CNN with a ResNet-50-FPN backbone is a sensible accuracy-oriented baseline. For low-latency inference on an edge device, a one-stage detector may be more appropriate. Benchmark on representative hardware rather than choosing from headline FPS figures.

Build a reliable dataset

Detection quality is usually limited by annotation quality and data coverage, not by a missing layer in the network. Use a consistent bounding-box convention: [xmin, ymin, xmax, ymax], with coordinates measured in pixels. Ensure xmax > xmin and ymax > ymin, and clip boxes to image boundaries.

A PyTorch detection dataset should return an image and a target dictionary containing:

boxes: a Float32 tensor shaped [N, 4];
labels: an Int64 tensor shaped [N];
image_id: a unique identifier;
area: box area, usually used by COCO-style evaluation;
iscrowd: normally zero unless your annotation format explicitly marks crowd regions.

Reserve label 0 for background. If your application has three foreground classes, num_classes must be 4.

Create train, validation, and test splits by location, camera, time period, or production batch, not only by randomly splitting near-duplicate frames. A random frame split can make metrics look excellent while hiding poor performance on a new Indian city, crop variety, factory shift, or camera. Keep a small “challenge set” containing difficult examples: tiny objects, glare, blur, occlusion, crowded scenes, and unusual weather.

Before training, run validation scripts that detect empty boxes, invalid coordinates, missing labels, duplicate annotations, unreadable images, and class IDs outside the configured range. Visualise random samples with boxes overlaid. This catches format errors faster than inspecting loss curves.

Use augmentations that preserve boxes

Object detection augmentations must modify the image and its bounding boxes together. Horizontal flips, moderate scale changes, crops, brightness shifts, contrast changes, and blur can improve robustness. Avoid transformations that create unrealistic examples or erase the visual evidence needed by small objects.

Albumentations is often practical because it handles common bounding-box transformations explicitly. Apply augmentations only to the training split. Validation and test images should represent the deployment distribution, with only deterministic resizing or normalisation.

For Indian deployments, test variation deliberately: harsh sunlight, monsoon rain, dust, low-light streets, regional vehicle types, different scripts on signboards, and compression from inexpensive cameras. Augmentation cannot replace missing real-world coverage, so track which conditions are represented and collect more data where errors cluster.

Initialise with transfer learning

Start from COCO-pretrained weights rather than training from scratch. Replace the classification head so that it predicts your classes:

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor


def get_model(num_classes):
    weights = torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights.DEFAULT
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=weights)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model

The weights API is preferable to older pretrained=True examples and is easier to maintain across modern torchvision releases. During early experiments, freeze part of the backbone if the dataset is small, then progressively unfreeze layers if validation performance plateaus. With a larger, visually distinct dataset, fine-tuning the full backbone is usually worthwhile.

Implement a stable training loop

Detection models in training mode return a loss dictionary. Faster R-CNN typically reports classification loss, box-regression loss, RPN objectness loss, and RPN box-regression loss. Sum these losses for backpropagation, but log each component separately; an abnormal RPN loss and an abnormal classifier loss suggest different problems.

Use a custom collate_fn because each image can contain a different number of boxes:

def collate_fn(batch):
    return tuple(zip(*batch))

A practical baseline is SGD with momentum, weight decay, and a scheduled learning rate. AdamW can also work, particularly during rapid experimentation, but compare optimisers using the same validation protocol. Use mixed precision on compatible GPUs to reduce memory use and improve throughput. Save checkpoints containing the model, optimiser, scheduler, epoch, class mapping, and configuration—not only the weights.

Track experiments with a fixed seed where possible, while remembering that GPU operations and data-loader behaviour can still introduce variation. Record the exact PyTorch, torchvision, CUDA, dataset version, and augmentation configuration. This matters when a model moves from a notebook to a grant-funded pilot or production service.

Evaluate with metrics that guide decisions

Accuracy is not an adequate detection metric. Use COCO-style mean average precision, including mAP@[.5:.95], AP at IoU 0.5, and separate results for small, medium, and large objects. Also report per-class precision, recall, and false positives per image or per hour of video.

IoU measures overlap between a predicted and ground-truth box. A model can have good AP at IoU 0.5 but poor localisation at stricter thresholds. Inspect precision-recall curves and choose confidence thresholds based on the business cost of misses versus false alerts. Never tune thresholds on the test set.

Break results down by conditions: camera, geography, time of day, weather, object size, and occlusion. A single aggregate score can conceal a model that performs well on cars but fails on auto-rickshaws, or works in daylight but collapses at night. Review false positives and false negatives manually every evaluation cycle.

Debug the failures before changing the architecture

Common failures have straightforward causes:

Loss does not fall: Check box format, label IDs, image scaling, target dtypes, and whether boxes are empty or invalid.
Predictions are all background: Confirm that the head was replaced and that num_classes includes background.
Small objects are missed: Increase input resolution, improve annotations, use a suitable feature pyramid, and collect more close-up examples.
Rare classes perform poorly: Add targeted data and balanced sampling; do not rely on weighted loss alone.
Validation is suspiciously strong: Look for duplicate scenes or adjacent video frames across splits.
Predictions are unstable across lighting: Add representative data, inspect normalisation, and calibrate confidence thresholds.

Class imbalance is often a data-collection problem. Oversampling rare examples can help, but it should not create repeated copies that encourage memorisation. Prefer more diverse examples and hard-negative mining from real deployment footage.

Prepare for deployment

A research checkpoint is not a product. Measure end-to-end latency, memory, image decoding time, post-processing, and power consumption on the target device. Export through TorchScript or ONNX where supported, and benchmark optimised runtimes such as TensorRT or OpenVINO. Validate exported outputs against PyTorch on a fixed image set; small differences are acceptable, changed class assignments are not.

For video, avoid running inference on every frame unless necessary. Use a suitable sampling rate, tracking between detections, batching where latency permits, and a clear policy for dropped frames. Add monitoring for confidence drift, new camera conditions, class distribution changes, and annotation-backed performance reviews. Store only the data needed for debugging and comply with applicable privacy and retention requirements.

Teams building several AI components should also plan their data and serving boundaries early. The same reproducibility principles used here complement custom neural network architecture design and the experiment discipline required for fine-tuning models on custom data, even though those topics use different model families.

A practical build checklist

Before calling the model production-ready, confirm that you have:

a documented class map and annotation policy;
validated boxes and visualised dataset samples;
splits that prevent location or frame leakage;
a baseline model and reproducible training configuration;
per-class and condition-specific metrics;
a reviewed false-positive and false-negative set;
an export and hardware benchmark;
monitoring, retraining, and rollback procedures.

For Indian founders, student teams, and research groups, this workflow turns PyTorch from a code sample into a measurable computer-vision system. Start with a narrow deployment target, prove performance on difficult local conditions, and expand classes only when the data and evaluation process can support them.

Apply for AI Grants India