Replies: 1 comment 6 replies
-
|
👋 Hello @sammlapp, thanks for opening this discussion and for your interest in training YOLO11 with a custom dataloader 🚀 This is an automated response to help you get started quickly — an Ultralytics engineer will also follow up on your question soon with more specific guidance. Since your topic relates to custom training and dataloaders, it’s especially helpful if you can share a bit more detail so we can assist effectively once an engineer joins the thread:
We also recommend a visit to the Docs 📚 where you can find many Python and CLI usage examples and where many of the most common questions may already be answered. If this is a 🐛 Bug Report related to your custom dataloader or training flow, please provide a minimum reproducible example (MRE) to help us debug it. If this is a custom training ❓ Question (as it appears), please provide as much information as possible, including:
Join the Ultralytics community where it suits you best:
UpgradeBefore diving deeper, please upgrade to the latest pip install -U ultralyticsEnvironmentsYOLO may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):
StatusIf this badge is green, all Ultralytics CI tests are currently passing. CI tests verify correct operation of all YOLO Modes and Tasks on macOS, Windows, and Ubuntu every 24 hours and on every commit. |
Beta Was this translation helpful? Give feedback.
-
|
Nice, that batch format is already very close to what YOLO11 expects; the main change is that You don’t need to re‑implement the loss; the import torch
from torch.utils.data import DataLoader
from ultralytics import YOLO
from ultralytics.models.yolo.detect import DetectionTrainer
class MyInMemoryDataset(torch.utils.data.Dataset):
def __getitem__(self, idx):
# return one sample as CHW tensor + per-image labels
img = ... # torch.Tensor [3, H, W]
cls = ... # (num_boxes,) long
bboxes = ... # (num_boxes, 4) xywhn
return {"img": img, "cls": cls, "bboxes": bboxes}
def __len__(self):
return ...
def collate_fn(batch):
imgs = torch.stack([b["img"] for b in batch], 0)
cls, bboxes, batch_idx = [], [], []
for i, b in enumerate(batch):
n = b["cls"].shape[0]
cls.append(b["cls"])
bboxes.append(b["bboxes"])
batch_idx.append(torch.full((n,), i, dtype=torch.long))
return {
"img": imgs,
"cls": torch.cat(cls),
"bboxes": torch.cat(bboxes),
"batch_idx": torch.cat(batch_idx),
}
class InMemoryTrainer(DetectionTrainer):
def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode="train"):
dataset = MyInMemoryDataset()
return DataLoader(dataset, batch_size=batch_size, shuffle=(mode == "train"), collate_fn=collate_fn)
model = YOLO("yolo11n.pt")
model.train(trainer=InMemoryTrainer, data="my_data.yaml", epochs=100)This way you keep the standard YOLO11 training loop and loss, and only swap out how batches are built; the pattern is the same as in the Advanced Customization guide and the |
Beta Was this translation helpful? Give feedback.
-
|
Awesome! subclassing DetectionTrainer and passing it to trainer argument was the piece I was missing. Thanks! |
Beta Was this translation helpful? Give feedback.
-
|
Glad it helped—if you later want to tweak loss, callbacks, or other hooks around your in‑memory dataset, the DetectionTrainer reference is a good place to see what you can override. |
Beta Was this translation helpful? Give feedback.
-
|
Thanks again for your help here. I've actually been encountering quite a few unexpected errors after implementing a Trainer subclass that overrides get_dataloader. I got some of them worked out, but am now stuck on an error that occurs in the Validator (which also expects paths to images). I can't pass So far my code looks like this, since internal code in the DetectionTrainer
from torch.utils.data import DataLoader
from ultralytics import YOLO
from ultralytics.models.yolo.detect import DetectionTrainer, DetectionValidator
from ultralytics.utils import LOGGER, RANK, nms, ops
import sys
sys.path.append("../detr")
from dataset import collate_to_yolo_batch, AudioToYoloDataset
from ultralytics.utils import DEFAULT_CFG
class CustomTrainer(DetectionTrainer):
def __init__(
self,
cfg=DEFAULT_CFG,
overrides=None,
_callbacks=None,
):
self.train_df = overrides.pop("train_df")
self.val_df = overrides.pop("val_df")
self.classes = overrides.get("classes")
self.preprocessor = overrides.pop("preprocessor")
super().__init__(cfg=cfg, overrides=overrides)
self.data = {
"nc": len(self.classes),
"channels": 3,
"names": self.classes,
"train": "notused",
"val": "notused",
}
def get_dataset(self):
pass
def plot_training_labels(self):
pass
def plot_training_samples(self, *args, **kwargs):
pass
def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode="train"):
if mode == "train":
dataset = AudioToYoloDataset(self.train_df, self.preprocessor)
dataset.preprocessor.bypass_augmentations = False
else:
dataset = AudioToYoloDataset(self.val_df, self.preprocessor)
dataset.preprocessor.bypass_augmentations = True
return DataLoader(
dataset,
batch_size=batch_size,
shuffle=(mode == "train"),
collate_fn=collate_to_yolo_batch,
)
class CustomValidator(DetectionValidator):
def _prepare_batch(self, si, batch):
"""Prepare a batch of images and annotations for validation.
Args:
si (int): Batch index.
batch (dict[str, Any]): Batch data containing images and annotations.
Returns:
(dict[str, Any]): Prepared batch with processed annotations.
"""
idx = batch["batch_idx"] == si
cls = batch["cls"][idx].squeeze(-1)
bbox = batch["bboxes"][idx]
ori_shape = batch["ori_shape"][si]
imgsz = batch["img"].shape[2:]
ratio_pad = batch["ratio_pad"][si]
if cls.shape[0]:
bbox = (
ops.xywh2xyxy(bbox)
* torch.tensor(imgsz, device=self.device)[[1, 0, 1, 0]]
) # target boxes
return {
"cls": cls,
"bboxes": bbox,
"ori_shape": ori_shape,
"imgsz": imgsz,
"ratio_pad": ratio_pad,
# "im_file": batch["im_file"][si],
}I make it through a training epoch, but then get an error in validation |
Beta Was this translation helpful? Give feedback.
-
|
You can’t pass from copy import copy
class CustomTrainer(DetectionTrainer):
def get_validator(self):
return CustomValidator(
self.test_loader,
save_dir=self.save_dir,
args=copy(self.args),
_callbacks=self.callbacks,
)Your current |
Beta Was this translation helpful? Give feedback.
Uh oh!
There was an error while loading. Please reload this page.
-
Hello, I would like to train YOLOv11 object detection with a custom dataloader. Rather than pre-computing a dataset of images stored on disk, my dataset and dataloader will perform in-memory preprocessing and pass the samples and bbox labels to the training step. Is this possible with the ultralytics.YOLO model object? Do I need to write the training loop from scratch like this example for yolo v5, and if so how should I compute the loss for the YOLO v11 object?
thanks!
Beta Was this translation helpful? Give feedback.
All reactions