How to integrate a custom dataloader for training?

[alevangel]

I have a custom PyTorch Dataloader, how can use it to train a new model? Where are the templates tu use?

PS this dataloader generates synthetic images on the fly.

[glenn-jocher]

@alevangel hello there! To integrate your custom PyTorch DataLoader for training with YOLOv8, you'll need to make sure that your DataLoader outputs the required data format expected by the YOLOv8 training pipeline. This typically includes images, targets, image paths, and shapes.

Once you have your DataLoader ready, you can modify the dataset loading part of the training script to use your custom DataLoader instead of the default one. Look for the section in the training script where datasets are instantiated and replace the default DataLoader with your own.

Keep in mind that your DataLoader should be compatible with the multi-processing environment if you plan to use distributed training. Also, ensure that your synthetic images and annotations are correctly formatted to match the expectations of the loss functions and metrics used during training.

For more detailed guidance on the expected data formats and integration steps, please refer to the documentation on custom datasets and training at https://docs.ultralytics.com. Good luck with your synthetic image generation, and happy training! 🚀

[lesept777]

Hi, are there any examples available that would make this process easier?

[halqadasi]

@pderrenger I tried to use it with data_loader parameter and it throws an error:

Traceback (most recent call last):

  File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3553, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)

  File ["<ipython-input-8-f70492c9994e>"](https://localhost:8080/#), line 5, in <cell line: 5>
    model.train(

  File "/usr/local/lib/python3.10/dist-packages/ultralytics/engine/model.py", line 799, in train
    self.trainer = (trainer or self._smart_load("trainer"))(overrides=args, _callbacks=self.callbacks)

  File "/usr/local/lib/python3.10/dist-packages/ultralytics/engine/trainer.py", line 101, in __init__
    self.args = get_cfg(cfg, overrides)

  File "/usr/local/lib/python3.10/dist-packages/ultralytics/cfg/__init__.py", line 294, in get_cfg
    check_dict_alignment(cfg, overrides)

  File "/usr/local/lib/python3.10/dist-packages/ultralytics/cfg/__init__.py", line 482, in check_dict_alignment
    raise SyntaxError(string + CLI_HELP_MSG) from e

  File "<string>", line unknown
SyntaxError: 'data_loader' is not a valid YOLO argument.

[JamesMcCullochDickens]

Is there any update to this?

[glenn-jocher]

Thank you for following up. Currently, the data_loader parameter is not supported directly in the model.train() method. To use a custom PyTorch DataLoader, you would need to implement a custom trainer by extending the base Trainer class. You can review the trainer implementation here for guidance. Let me know if you need further clarification!

[JamesMcCullochDickens]

@glenn-jocher before I dive into all this code, which is quite complex, is there a simple example lying around somewhere for pose models with a custom dataloader that subclasses from the PoseTrainer class? I have a bit of an odd use case so I really need this custom dataloader.

Also is there documentation anywhere for what the input batch dictionary expects for pose models?
Supposing I have a batch of N images,

I'm imagining:
batch["im"] is a tensor of N images in (N, 3, H, W) form
batch["keypoints"] is a list of keypoint tensors, where each entry is of the form (num_boxes, 17, 3)
batch["boxes"] is a list of bounding box tensors in yolo format normalized (x_c, y_c, w, h) and
batch["batch_idx"] is a list of tensors where for each entry you list the batch value per bounding box

[glenn-jocher]

Thank you for your question! Unfortunately, there isn’t a simple example readily available for subclassing PoseTrainer with a custom dataloader. However, you can refer to the PoseTrainer class implementation here for guidance on customizing it.

For your second query, your assumptions about the input batch dictionary are mostly correct. The expected format for pose models is as follows:

• batch["im"]: Tensor of N images, shape (N, 3, H, W).
• batch["keypoints"]: Tensor of shape (num_boxes, num_keypoints, 3), where 3 represents [x, y, visible].
• batch["bboxes"]: Tensor of bounding boxes in YOLO format (x_c, y_c, w, h), normalized.
• batch["batch_idx"]: Tensor indicating the batch index for each bounding box.

Feel free to share any follow-up questions or clarifications!