Custom model training

[someuser2026]

Hi! Thank you for the great library. It is much easier to work with compared to others.

The built in yolo models and configs work great, however, i wanted to try with some different and custom models including swin, deformable detr and more niche ones. I tried to go through the docs and code, howeer was not fully able to understand the strcuture expected of new backbones or necks for optimal integration.

Would greatly appreciate if you could help me with:

1. How does the config files and parameters work - what is order in which they are supplied to the underlying functions or is there no order at all?
2. What combinations of blocks or modules are supported? Say for example can i simply replace the backbone for yolo11 with transformer blocks and it would work?
3. When writing custom backbones what strucutre of inputs and outputs is expected to integrate with the rest of the pipeline. Especially the loss and metric calculations
4. Building upon the previous question, for custom models does the config also needs to have some specific structure?

Sorry if these are repeats. I tried my best to read through the discussions an only include parts that are not explained.
Thank you. Greatly appreciate your time.
Sincerely

[UltralyticsAssistant]

👋 Hello @someuser2026, thank you for your kind words and interest in Ultralytics 🚀! We’re glad to hear you’ve had a positive experience so far.

If this is a ❓ custom model integration or training question, please provide as much information as possible about your use case, including any config files, code snippets, or error messages. If you encounter any issues or bugs, please include a minimum reproducible example to help us assist you effectively.

For your questions about integrating custom backbones or modules, we recommend reviewing the Docs—especially the sections on Python and CLI usage, as well as our Tips for Best Training Results. You can also check out existing model definitions in the repository to see how backbones and necks are structured.

Join the Ultralytics community wherever you prefer! For real-time chat, visit Discord 🎧. For more in-depth discussions, check out Discourse. Or connect with others on our Subreddit.

Upgrade

Please make sure you are using the latest ultralytics package and all requirements in a Python>=3.8 environment with PyTorch>=1.8:

pip install -U ultralytics

Environments

Run YOLO in any of these up-to-date verified environments (with dependencies including CUDA/CUDNN, Python, and PyTorch preinstalled):

• Notebooks with free GPU:
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Amazon Deep Learning AMI. See AWS Quickstart Guide
• Docker Image. See Docker Quickstart Guide

Status

If 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.

This is an automated response 🤖. An Ultralytics engineer will review your questions and provide more detailed guidance soon!

[VeeranjaneyuluToka]

@someuser2026 , Some of the above functionalities are supported in mmyolo or mmdetection in case if you want to explore.

[someuser2026]

Thank you. Appreciate the help.

I did try but installing that stack is very difficult. I was not able to install it in a clean env. Are there any nice guides that i am unaware of?

Thank you!

[VeeranjaneyuluToka]

@someuser2026 , I generally use their docker images and they work almost perfect.

[someuser2026]

Ok. Thank you! I want to use it on an hpc cluster and docker is not an option. Appreciate the help!

[Y-T-G]

You can check the explanations here to understand the YAML configs:

#19615
#18564 (comment)

[someuser2026]

Hi! Thank you! Your upgrade and post are very helpful! One more question: when we use the torchvision module to use a predefined backbone and use the arch for training, is the backbone also trainable or are the backbone weights always frozen?

[Y-T-G]

It's trainable

[someuser2026]

ok. Thank you!

[someuser2026]

Hi, tired adding the custom module using the above method but is says KeyError. It seems it is not being registered. Has the custom module function merged into the current lib? Attached are the error trace and the yaml file.

yaml file

# 0 torch.Size([1, 3, 640, 640])
# 1 torch.Size([1, 160, 160, 96])
# 2 torch.Size([1, 160, 160, 96])
# 3 torch.Size([1, 80, 80, 192])
# 4 torch.Size([1, 80, 80, 192])
# 5 torch.Size([1, 40, 40, 384])
# 6 torch.Size([1, 40, 40, 384])
# 7 torch.Size([1, 20, 20, 768])
# 8 torch.Size([1, 20, 20, 768])
# 9 torch.Size([1, 20, 20, 768])
# 10 torch.Size([1, 768, 20, 20])
# 11 torch.Size([1, 768, 1, 1])


nc: 1
scales:
  x: [1.00, 1.00, 1024]
backbone:
  - [-1, 1, TorchVision, [768, "swin_t", "DEFAULT", True, 2, True]]  # - 0
head:
  - [0, 1, Index, [96, 2]]  # extracts 2th output (1, 192, 80, 80) - 1
  - [0, 1, Index, [192, 4]]  # extracts 4th output (1, 192, 80, 80) - 2
  - [0, 1, Index, [384, 6]]  # extracts 6th output (1, 384, 40, 40) - 3
  - [0, 1, Index, [768, 9]]  # extracts 8th output (1, 768, 20, 20) - 4
  
  # Add a custom permutation step after each Index extraction
  - [1, 1, Permute, [[0, 3, 1, 2]]]  # Convert (B,H,W,C) to (B,C,H,W)
  - [2, 1, Permute, [[0, 3, 1, 2]]]
  - [3, 1, Permute, [[0, 3, 1, 2]]]
  - [4, 1, Permute, [[0, 3, 1, 2]]]

  - [4, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 3], 1, Concat, [1]] # cat backbone P4
  - [-1, 2, C3k2, [512, False]] # 7

  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 2], 1, Concat, [1]] # cat backbone P3
  - [-1, 2, C3k2, [256, False]] # 10 (P3/8-small)

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 7], 1, Concat, [1]] # cat head P4
  - [-1, 2, C3k2, [512, False]] # 13 (P4/16-medium)

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 4], 1, Concat, [1]] # cat head P5
  - [-1, 2, C3k2, [1024, True]] # 16 (P5/32-large)

  - [[10, 13, 16], 1, OBB, [nc, 1]] # pass all three feature to obb head
  # - [[1, 2, 3], 1, OBB, [nc, 1]] # pass all three feature to obb head

module: 
   init: |
     # Define all custom modules here.
     # Make sure the indentations are correct.

      import torch.nn as nn
      import torch

      class Permute(nn.Module):
          def __init__(self, permutation: list[int] | tuple[int]):
              assert all(isinstance(p, int) for p in permutation), "All permutation indices should be integers"
              self.permutation = permutation

          def forward(self, x: torch.Tensor):
              assert x.dim() == len(self.permutation)
              return torch.permute(x, self.permutation)

parse: |
  # Parser for args.
  if m is Permute:
    c1 = ch[f]

Error trace

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[65], line 1
----> 1 model = YOLO("../../configs/ultralytics/swin/swin_t-yolo11x-obb.yaml", task="detect")
      2 img = torch.randn(1, 3, 640, 640).clip(0, 1)
      3 img

File ~/miniforge3/envs/ultralytics/lib/python3.13/site-packages/ultralytics/models/yolo/model.py:79, in YOLO.__init__(self, model, task, verbose)
     76     self.__dict__ = new_instance.__dict__
     77 else:
     78     # Continue with default YOLO initialization
---> 79     super().__init__(model=model, task=task, verbose=verbose)
     80     if hasattr(self.model, "model") and "RTDETR" in self.model.model[-1]._get_name():  # if RTDETR head
     81         from ultralytics import RTDETR

File ~/miniforge3/envs/ultralytics/lib/python3.13/site-packages/ultralytics/engine/model.py:149, in Model.__init__(self, model, task, verbose)
    147 __import__("os").environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"  # to avoid deterministic warnings
    148 if str(model).endswith((".yaml", ".yml")):
--> 149     self._new(model, task=task, verbose=verbose)
    150 else:
    151     self._load(model, task=task)

File ~/miniforge3/envs/ultralytics/lib/python3.13/site-packages/ultralytics/engine/model.py:261, in Model._new(self, cfg, task, model, verbose)
    259 self.cfg = cfg
    260 self.task = task or guess_model_task(cfg_dict)
--> 261 self.model = (model or self._smart_load("model"))(cfg_dict, verbose=verbose and RANK == -1)  # build model
    262 self.overrides["model"] = self.cfg
    263 self.overrides["task"] = self.task

File ~/miniforge3/envs/ultralytics/lib/python3.13/site-packages/ultralytics/nn/tasks.py:397, in DetectionModel.__init__(self, cfg, ch, nc, verbose)
    395     LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
    396     self.yaml["nc"] = nc  # override YAML value
--> 397 self.model, self.save = parse_model(deepcopy(self.yaml), ch=ch, verbose=verbose)  # model, savelist
    398 self.names = {i: f"{i}" for i in range(self.yaml["nc"])}  # default names dict
    399 self.inplace = self.yaml.get("inplace", True)

File ~/miniforge3/envs/ultralytics/lib/python3.13/site-packages/ultralytics/nn/tasks.py:1671, in parse_model(d, ch, verbose)
   1646 repeat_modules = frozenset(  # modules with 'repeat' arguments
   1647     {
   1648         BottleneckCSP,
   (...)   1663     }
   1664 )
   1665 for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]):  # from, number, module, args
   1666     m = (
   1667         getattr(torch.nn, m[3:])
   1668         if "nn." in m
   1669         else getattr(__import__("torchvision").ops, m[16:])
   1670         if "torchvision.ops." in m
-> 1671         else globals()[m]
   1672     )  # get module
   1673     for j, a in enumerate(args):
   1674         if isinstance(a, str):

KeyError: 'Permute'

[Y-T-G]

It's not merged. You will need to manually install the PR.

pip install git+https://github.com/ultralytics/ultralytics@custom_module

[WangCarl77424]

Hello, I have a question.
When I do custom training.
My training folders are named image and label.
The image folder contains image and txt.
The label folder can only be trained if it has the same txt.

When my training folders are named images and labels with s
The image folder only has image
The label folder can be trained if it has txt

The yaml files are all correct
I would like to ask why the image folder needs to contain image and txt when s is not added

What will be the impact of this?

[Y-T-G]

The correct format is images and labels because Ultralytics searches for all the images based on file extension in subfolders and then performs string replacement of images in file path to labels to get labels.

images/1.jpg becomes labels/1.txt

If you use any other name, like image then the replacement will not work. So it is still looking in the same directory for labels.

[someuser2026]

Hi, I tried using a swin_t backbone with yolo11x neack and obb head. However, it appears that the tensors returned by the swin bakcbone have channels in the last dimension wheras the yolo layers accpet it in 2nd dimension. Is there a layer to help reshape the input? Attached is the current yaml file:

# 0 torch.Size([1, 3, 640, 640])
# 1 torch.Size([1, 160, 160, 96])
# 2 torch.Size([1, 160, 160, 96])
# 3 torch.Size([1, 80, 80, 192])
# 4 torch.Size([1, 80, 80, 192])
# 5 torch.Size([1, 40, 40, 384])
# 6 torch.Size([1, 40, 40, 384])
# 7 torch.Size([1, 20, 20, 768])
# 8 torch.Size([1, 20, 20, 768])
# 9 torch.Size([1, 20, 20, 768])
# 10 torch.Size([1, 768, 20, 20])
# 11 torch.Size([1, 768, 1, 1])


nc: 1
scales:
  x: [1.00, 1.00, 1024]
backbone:
  - [-1, 1, TorchVision, [768, "swin_t", "DEFAULT", True, 2, True]]  # - 0
head:
  - [0, 1, Index, [96, 2]]  # extracts 2th output (1, 192, 80, 80) - 1
  - [0, 1, Index, [192, 4]]  # extracts 4th output (1, 192, 80, 80) - 2
  - [0, 1, Index, [384, 6]]  # extracts 6th output (1, 384, 40, 40) - 3
  - [0, 1, Index, [768, 9]]  # extracts 8th output (1, 768, 20, 20) - 4

  - [4, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 3], 1, Concat, [1]] # cat backbone P4
  - [-1, 2, C3k2, [512, False]] # 7

  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 2], 1, Concat, [1]] # cat backbone P3
  - [-1, 2, C3k2, [256, False]] # 10 (P3/8-small)

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 7], 1, Concat, [1]] # cat head P4
  - [-1, 2, C3k2, [512, False]] # 13 (P4/16-medium)

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 4], 1, Concat, [1]] # cat head P5
  - [-1, 2, C3k2, [1024, True]] # 16 (P5/32-large)

  - [[10, 13, 16], 1, OBB, [nc, 1]] # pass all three feature to obb head

[glenn-jocher]

The issue you're encountering is that Swin Transformer outputs tensors in (B, H, W, C) format while YOLO expects (B, C, H, W) format. You'll need to add a permutation step to transpose the dimensions.

You can create a custom permutation module or use a workaround. Here's a solution using a custom module approach - add this to your YAML after the Index extractions:

# Add a custom permutation step after each Index extraction
- [1, 1, lambda, ['x: x.permute(0, 3, 1, 2)']]  # Convert (B,H,W,C) to (B,C,H,W)
- [2, 1, lambda, ['x: x.permute(0, 3, 1, 2)']]  
- [3, 1, lambda, ['x: x.permute(0, 3, 1, 2)']]
- [4, 1, lambda, ['x: x.permute(0, 3, 1, 2)']]

However, lambda functions aren't directly supported in YAML configs. A cleaner approach would be to modify the Index module to handle this automatically, or create a custom module that combines index extraction with permutation.

Alternatively, you could subclass the existing modules to handle the dimension conversion. The core issue is that vision transformers typically output in (H, W, C) format while CNNs expect (C, H, W) format, so you'll need this conversion step whenever mixing these architectures.

[someuser2026]

Hi, Thank you! Appreciate the help. Do you think that having a reshape or maybe even a 'lambda' module would be nice in the library or would it just bloat id uncessarliy since the library mainly focuses on the YOLO and CNN models? Just asking. Though i gues having a "lambda" module directly might lead to security issues.

[glenn-jocher]

You're absolutely right about the security concerns with lambda functions - executing arbitrary code from YAML configs would be a significant security risk.

A simple Permute or Reshape module would actually be quite valuable without bloating the library, especially as more users integrate transformer backbones with YOLO architectures. These are fundamental tensor operations that bridge different model families, so they align well with Ultralytics' goal of supporting diverse architectures.

The library already includes utility modules like Concat and Index, so adding basic tensor manipulation modules like Permute would be consistent with that approach. It's a clean, secure way to handle the common (B,H,W,C) to (B,C,H,W) conversion needed when mixing vision transformers with CNN-based heads.

This kind of flexibility is becoming increasingly important as the field moves toward hybrid architectures, so it would be a practical addition rather than unnecessary bloat.

[someuser2026]

ok. Thank you!

[glenn-jocher]

You're welcome! Feel free to reach out if you have any other questions as you continue working with custom architectures. Good luck with your Swin Transformer + YOLO11 integration! 🚀