Training with a custom (dynamic) dataloader/dataset.

[vshas]

I want to use a dataloader/dataset that generates training examples on the fly. As opposed to the standard (static) situation of having directories with images and masks/labels. In my specific example I am considering a segmentation task where I have several, big, images with the corresponding masks. As a means of data augmentation I want to sample/cut from each big image/mask pair, a smaller image/mask pair to use as training example. In this way I can augment my data and it is also easier to label medium amount of big images vs a high amount of small images. Now I have tried to train the 'yolov8n-seg.pt' model using a custom Dataset (torch.utils.data) and accompanying DataLoader. But even after quite some time hacking, I cant get it to work. Here is my code:

from torch.utils.data import Dataset, DataLoader
import cv2
import numpy as np
import random
from ultralytics import YOLO
from ultralytics.models.yolo.segment.train import SegmentationTrainer

class CustomYOLODataset(Dataset):
    def __init__(self, base_path, input_size=1440, output_size=480, num_samples=100):
        self.base_path = base_path
        self.output_size = output_size
        self.input_size = input_size
        self.image_paths = os.listdir(base_path+'/images')
        self.image_annotations = {}
        self.num_samples = num_samples
        for file in os.listdir(base_path+'/annotations'):
            with open(base_path+'/annotations/' + file) as f:
                self.image_annotations.update(json.load(f))
        self.image_selection = [random.randint(0, len(self.image_paths)-1) for _ in range(self.num_samples)]
        self.crop_coordinates = [(x, y) for (x, y) in zip(
                [random.randint(0, self.input_size - self.output_size) for _ in range(self.num_samples)],
                [random.randint(0, self.input_size - self.output_size) for _ in range(self.num_samples)],
            )
        ]

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        img_path = self.image_paths[self.image_selection[idx]]
        img = cv2.imread(self.base_path + '/images/' + img_path)
        mask = np.zeros((1440, 1440), dtype=np.uint8)
        for ann in self.image_annotations[img_path]:
            polygon_points = np.array(ann[1:], dtype=np.int32).reshape(-1, 2)  # Convert to Nx2 array
            cv2.fillPoly(mask, [polygon_points], color=1)

        x1, y1 = self.crop_coordinates[idx]
        img_crop = img[y1:y1+self.output_size, x1:x1+self.output_size]
        mask_crop = mask[y1:y1+self.output_size, x1:x1+self.output_size]
        img_tensor = torch.tensor(img_crop, dtype=torch.float32).permute(2, 0, 1) / 255.0
        mask_tensor = torch.tensor(mask_crop, dtype=torch.long)
        return img_tensor, mask_tensor


class MyCustomTrainer(SegmentationTrainer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def get_dataset(self):
        data = {}
        data["nc"] = 1
        data["names"] = ["class1"]
        self.data = data
        return 'segmentation_set_test/train', 'segmentation_set_test/test'

    def get_dataloader(self, base_path, batch_size, **kwargs):
        dataset =CustomYOLODataset(base_path, input_size=1440, output_size=480)
        dataloader = DataLoader(dataset, batch_size=2, shuffle=False)
        return dataloader

    def train(self):
        super().train()

model = YOLO('yolov8n-seg.pt')

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

trainer = MyCustomTrainer
#
model.train(epochs=2, trainer=trainer, resume=True)

You can clearly see my hacking attempts in the get_dataset method to circumvent the dataset.yaml. I also tried it with dataset.yaml but I could not get it to work nonetheless. Can somebody explain to me what I am doing wrong? Or is this kind of thing not supported? I do not believe I am doing very exotic stuff here. Thanks in advance!

[UltralyticsAssistant]

👋 Hello @vshas, thank you for sharing your detailed use case and for your interest in Ultralytics 🚀! Your implementation of a custom dataset and dynamic data loading setup for YOLOv8 segmentation tasks is fascinating, and I'm sure the community will greatly benefit from this discussion. 😊

We understand the challenges you’ve faced in integrating your custom dataset into the YOLOv8 training pipeline. Although this particular setup may require some extra steps, it seems like it can definitely work with the right modifications. We'll do our best to help guide you!

If this is a 🐛 Bug Report—i.e., something appears broken—please ensure that the issue is reproducible and provide a minimum reproducible example (MRE) that isolates the problem. It allows us to better identify any potential bugs within the training process or framework integration. Your code snippet offers a great starting point, so thank you for sharing that! 🙌

If this is a ❓ Custom Training Inquiry, here are a few suggestions to further refine your setup:

1. Ensure you have reviewed our Tips for Best Training Results. This guide is particularly helpful for understanding dataset preparation, augmentation, and training logic.
2. Double-check that your Yolov8 custom trainer class aligns with the framework’s structure. The get_dataset and get_dataloader methods might require proper inputs, testing, and output formats specific to Ultralytics YOLOv8.
3. Explore the Docs for more details about how YOLOv8 handles data loading and custom configurations, especially for tasks like segmentation.
4. Validate that the directory and file structures of your input dataset are complete and properly referenced in your code if using the dataset.yaml workflow.

Additionally, to debug effectively and streamline experimentation:

• Consider upgrading to the latest ultralytics package to ensure you’re working with the most recent and stable features:

pip install -U ultralytics

• Test smaller slices of the pipeline (e.g., just initialize and test the CustomYOLODataset and DataLoader) before integrating.

If you’re interested in community exchange, you can also join us on Discord 🎧 for real-time discussions or post your progress and questions on our Discourse or Subreddit!

Lastly, stay tuned for assistance from an Ultralytics engineer who will take a closer look at your code and setup soon. Thank you again for your patience and for diving into the YOLOv8 framework—your efforts contribute greatly to the improvement of the ecosystem. 🚀✨

[glenn-jocher]

@vshas your custom dataset must output images and labels exactly as expected by our segmentation pipeline—ensure your dynamic cropping, transforms, and annotations align with YOLO's required format.

[vshas]

For anyone with a similar issue please also see #10174. At that thread there is some back and forward but I think the answer is that you have to figure it out yourself. Basically get deep into the ultralytics source code and make it work.

[glenn-jocher]

@vshas for custom dataset integration, we recommend subclassing YOLODataset from ultralytics/data/dataset.py instead of torch.utils.data.Dataset to maintain compatibility with YOLO's data pipeline. Your dataset needs to return properly formatted segmentation masks with associated instance boundaries and class IDs. The get_dataloader override should match the expected signature from ultralytics/engine/trainer.py, including handling mode and rank parameters. For dynamic cropping, implement your logic in the dataset's __getitem__ while maintaining YOLO's required output format (bounding boxes + segmentation masks). You'll need to ensure proper data configuration through either a YAML file or programmatic setup in your trainer class.