Attention mechanisms in yolov8, yolov9

[AnirudhSanthana]

Hello Ultralytics Team,

I’m trying to add attention mechanism to yolov8 and yolov9. I’ve already defined it in the YAML file but I’m unsure of the next steps. Additionally I noticed there’s no ' common.py' file to add the attention block. Could you please outline the exact steps and specify which files need to be modified after incorporating the block in the YAML file? @glenn-jocher

Thank you

Best regards,
Anirudh

[UltralyticsAssistant]

👋 Hello @AnirudhSanthana, thank you for your interest in Ultralytics 🚀! We recommend checking out our Documentation, which provides extensive resources and examples for using and customizing YOLO models. Attention mechanisms are an exciting area of customization, and your project sounds intriguing! 🔍

If this is a 🐛 Bug Report, please ensure you provide a minimum reproducible example (MRE) so we can better understand and help resolve any issues you're encountering.

If this is a custom ❓ Question about integrating attention mechanisms, we'd love for you to provide as much detail as possible to assist our engineers and the community in supporting your inquiry. This could include your YAML configurations, any modifications you've started, or errors you're encountering. Additionally, you might find inspiration in our Tips for Best Training Results guide.

Join the Ultralytics community where it suits you best! For real-time tips, jump into our Discord 🎧. For longer discussions with other enthusiasts and experts, visit Discourse. Or engage on our Subreddit to share progress and ideas.

Upgrade

Make sure you're using the latest ultralytics package to take advantage of all the latest features and updates. Upgrade in a Python>=3.8 environment with PyTorch>=1.8:

pip install -U ultralytics

Environments

To experiment with customizations like attention mechanisms, be sure to use a verified, up-to-date environment. Options include:

• 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. These tests validate the operation of all YOLO Modes and Tasks across multiple platforms (macOS, Windows, Ubuntu) every 24 hours and after every commit.

This is an automated response to ensure you get the most relevant information quickly 😊. An Ultralytics engineer will be available soon to assist further!

[glenn-jocher]

@AnirudhSanthana hello Anirudh,

To incorporate an attention mechanism into YOLOv8 or YOLOv9, you need to define the custom block in a relevant module file, such as ultralytics/nn/modules/block.py. After defining the block, ensure its integration in the parse_model function located in ultralytics/nn/tasks.py. This function processes the YAML file and maps the block to its implementation.

For detailed guidance, refer to the block module documentation and ensure consistency across the YAML and Python files. If you encounter specific issues, feel free to share them for support.

[glenn-jocher]

It seems like you've thoroughly implemented the CBAM integration, but the KeyError: 'CBAM' suggests the module is not being correctly registered. A few things to verify:

1. Ensure the CBAM class is correctly defined and resides in ultralytics/nn/modules/block.py.
2. Double-check that CBAM is imported correctly in tasks.py using from ultralytics.nn.modules.block import CBAM.
3. Confirm that CBAM is included in both the base_modules and repeat_modules sets in parse_model.
4. Verify that the yolov8.yaml file references CBAM with the same name and argument format as defined in tasks.py.

If these points are correct and the issue persists, try debugging with print(globals()) in parse_model to confirm that CBAM is accessible in the global scope. Let me know if you need further assistance!

[AnirudhSanthana]

Hello Glenn @glenn-jocher

I'm reaching out again regarding my attempt to add the CBAM module to my YOLOv8 setup. I’ve attached my YAML, tasks.py, and block.py files once more, with clear comments marking all the changes I’ve made for CBAM integration. I’ve tried countless times to debug but I’m truly stuck....

Would you be able to take another look and let me know where the issue might be? I’m sorry for the repeated requests—I’ve spent so much time troubleshooting and I’m at a complete standstill... Any help or guidance you can offer would be much appreciated!

Thank you again for your time and support.
block.txt
tasks.txt
yaml.txt

[glenn-jocher]

Thank you for sharing your files and providing detailed context. The issue might stem from how the CBAM module is being registered or referenced. Please confirm that the CBAM class is accessible in the global scope within tasks.py by printing globals() and checking for its inclusion. If the issue persists, carefully review for any typos or mismatches in the YAML file, as module names are case-sensitive. Let us know if further clarification is needed!

[AnirudhSanthana]

Thank you for your previous guidance @glenn-jocher. While following your suggestions, I encountered a new issue: KeyError: 'CBAM'. Despite ensuring that CBAM is properly defined, registered, and referenced, the error persists. I believe the issue stems from the parse_model function in tasks.py and the CBAM definition in block.py. I’ve included both below for reference:

parse_model Function in tasks.py

def parse_model(d, ch, verbose=True): # model_dict, input_channels(3)
"""
Parse a YOLO model.yaml dictionary into a PyTorch model.

Args:
    d (dict): Model definition loaded from YAML.
    ch (int): Number of input channels.
    verbose (bool): If True, print the module information.

Returns:
    (nn.Sequential, list): The constructed model and a list of layer indices whose outputs are saved.
"""
import ast

# Args
legacy = True  # backward compatibility for v3/v5/v8/v9 models
max_channels = float("inf")
nc, act, scales = (d.get(x) for x in ("nc", "activation", "scales"))
depth, width, kpt_shape = (d.get(x, 1.0) for x in ("depth_multiple", "width_multiple", "kpt_shape"))
if scales:
    scale = d.get("scale")
    if not scale:
        scale = tuple(scales.keys())[0]
        LOGGER.warning(f"WARNING ⚠️ no model scale passed. Assuming scale='{scale}'.")
    depth, width, max_channels = scales[scale]

# Override default activation if present in YAML
if act:
    Conv.default_act = eval(act)  # e.g. Conv.default_act = nn.SiLU()
    if verbose:
        LOGGER.info(f"{colorstr('activation:')} {act}")

# Logging header
if verbose:
    LOGGER.info(f"\n{'':>3}{'from':>20}{'n':>3}{'params':>10}  {'module':<45}{'arguments':<30}")

ch = [ch]  # list of output channel dimensions; starts with 'ch' input channels
layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch out

# ----------------------------------------------------------------------------------------
# 1) Register YOLO base modules (including CBAM). 
# ----------------------------------------------------------------------------------------
base_modules = frozenset({
    Classify,
    Conv,
    ConvTranspose,
    GhostConv,
    Bottleneck,
    GhostBottleneck,
    SPP,
    SPPF,
    C2fPSA,
    C2PSA,
    DWConv,
    Focus,
    BottleneckCSP,
    C1,
    C2,
    C2f,
    C3k2,
    RepNCSPELAN4,
    ELAN1,
    ADown,
    AConv,
    SPPELAN,
    C2fAttn,
    C3,
    C3TR,
    C3Ghost,
    nn.ConvTranspose2d,
    DWConvTranspose2d,
    C3x,
    RepC3,
    PSA,
    SCDown,
    C2fCIB,
    CBAM,  # <--- CBAM ADDED HERE
})

# If you want CBAM repeated multiple times, add it below. Otherwise omit.
repeat_modules = frozenset({
    BottleneckCSP,
    C1,
    C2,
    C2f,
    C3k2,
    C2fAttn,
    C3,
    C3TR,
    C3Ghost,
    C3x,
    RepC3,
    C2fPSA,
    C2fCIB,
    C2PSA,
    CBAM,
})

# ----------------------------------------------------------------------------------------
# 2) Parse model backbone + head layers from YAML
# ----------------------------------------------------------------------------------------
for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]):  # from, number, module, args
    # Dynamically resolve the module 'm' from the string in YAML
    m = (
        getattr(torch.nn, m[3:]) if "nn." in m
        else getattr(__import__("torchvision").ops, m[16:]) if "torchvision.ops." in m
        else globals()[m]
    )

    # Safely evaluate any string arguments (e.g. "64", "True", or "[1,2]")
    for j, a in enumerate(args):
        if isinstance(a, str):
            with contextlib.suppress(ValueError):
                args[j] = locals()[a] if a in locals() else ast.literal_eval(a)

    n = n_ = max(round(n * depth), 1) if n > 1 else n  # depth gain

    # ------------------------------------------------------------------------------------
    # 2A) If the module is in 'base_modules'
    # ------------------------------------------------------------------------------------
    if m in base_modules:
        # ----------------------
        # Special case: CBAM
        # ----------------------
        if m is CBAM:
            # forcibly set inbound channels = ch[f]
            c1 = ch[f]
            # override the first argument if needed
            if len(args) >= 1:
                args[0] = c1
            else:
                args.insert(0, c1)

            # CBAM is in-place, so output channels = input channels
            c2 = c1

            # If also in 'repeat_modules'
            if m in repeat_modules:
                # Insert repeat count at index 2 (c1, c2, then 'n')
                args.insert(2, n)
                n = 1

        # -------------------------
        # General YOLO modules
        # -------------------------
        else:
            c1, c2 = ch[f], args[0]
            if c2 != nc:  # if c2 not equal to number of classes for Classify() etc.
                c2 = make_divisible(min(c2, max_channels) * width, 8)

            # Example special logic for C2fAttn
            if m is C2fAttn:
                args[1] = make_divisible(min(args[1], max_channels // 2) * width, 8)
                args[2] = int(
                    max(round(min(args[2], max_channels // 2 // 32)) * width, 1) if args[2] > 1 else args[2]
                )

            # Rebuild args: [c1, c2, *args[1:]]
            args = [c1, c2, *args[1:]]

            # If the module is also in 'repeat_modules', handle repeats
            if m in repeat_modules:
                args.insert(2, n)  # number of repeats after c1, c2
                n = 1

            # For C3k2, optional legacy toggles
            if m is C3k2:
                legacy = False
                if scale in "mlx":
                    args[3] = True

    # ------------------------------------------------------------------------------------
    # 2B) Other special modules
    # ------------------------------------------------------------------------------------
    elif m is AIFI:
        args = [ch[f], *args]

    elif m in frozenset({HGStem, HGBlock}):
        c1, cm, c2 = ch[f], args[0], args[1]
        args = [c1, cm, c2, *args[2:]]
        if m is HGBlock:
            args.insert(4, n)  # number of repeats
            n = 1

    elif m is ResNetLayer:
        c2 = args[1] if args[3] else args[1] * 4

    elif m is nn.BatchNorm2d:
        args = [ch[f]]

    elif m is Concat:
        c2 = sum(ch[x] for x in f)

    elif m in frozenset({Detect, WorldDetect, Segment, Pose, OBB, ImagePoolingAttn, v10Detect}):
        args.append([ch[x] for x in f])  # pass in channel list
        if m is Segment:
            args[2] = make_divisible(min(args[2], max_channels) * width, 8)
        if m in {Detect, Segment, Pose, OBB}:
            m.legacy = legacy

    elif m is RTDETRDecoder:  # channels arg must be inserted at index 1
        args.insert(1, [ch[x] for x in f])

    elif m is CBLinear:
        c2 = args[0]
        c1 = ch[f]
        args = [c1, c2, *args[1:]]

    elif m is CBFuse:
        c2 = ch[f[-1]]

    elif m in frozenset({TorchVision, Index}):
        c2 = args[0]
        c1 = ch[f]
        args = [*args[1:]]

    else:
        # default fallback if module not in base_modules
        c2 = ch[f]

    # ------------------------------------------------------------------------------------
    # 3) Build the module(s). If n>1, replicate it in nn.Sequential
    # ------------------------------------------------------------------------------------
    m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)

    # ------------------------------------------------------------------------------------
    # 4) Attach meta info to the module and log
    # ------------------------------------------------------------------------------------
    t = str(m)[8:-2].replace("__main__.", "")  # module type (string)
    m_.np = sum(x.numel() for x in m_.parameters())  # number params
    m_.i, m_.f, m_.type = i, f, t  # index, 'from', type
    if verbose:
        LOGGER.info(f"{i:>3}{str(f):>20}{n_:>3}{m_.np:10.0f}  {t:<45}{str(args):<30}")

    # ------------------------------------------------------------------------------------
    # 5) Update save list and channel list
    # ------------------------------------------------------------------------------------
    save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # track layers to save
    layers.append(m_)
    if i == 0:
        ch = []  # reset channel list if this is the first layer
    ch.append(c2)

return nn.Sequential(*layers), sorted(save)

CBAM Definition in block.py

class CBAM(nn.Module):
def init(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max'], no_spatial=False):
super(CBAM, self).init()
self.ChannelGate = ChannelGate(gate_channels, reduction_ratio, pool_types)
self.no_spatial = no_spatial
if not no_spatial:
self.SpatialGate = SpatialGate()

def forward(self, x):
    x_out = self.ChannelGate(x)
    if not self.no_spatial:
        x_out = self.SpatialGate(x_out)
    return x_out

This error is new and I’ve attached an image of the traceback for clarity. Could you kindly review this and let me know what I might have missed or done incorrectly?

Thank you so much for your time and support!

[glenn-jocher]

The KeyError: 'CBAM' indicates that the CBAM module is not being recognized in the global scope within the parse_model function. Ensure the following:

1. Importing CBAM: Confirm that CBAM is correctly imported in tasks.py using:

   from ultralytics.nn.modules.block import CBAM

2. Global Scope: Add this debug check in the parse_model function to verify if CBAM is listed:

   print(globals().keys())

   Make sure CBAM appears in the output.

3. YAML Reference: Double-check the YAML file to ensure CBAM is spelled correctly and matches the class definition. Module names are case-sensitive.

If the issue persists, the traceback image suggests the failure occurs during the instantiation of CBAM. This might also indicate a mismatch in expected arguments. Verify that the arguments passed in the YAML file match the CBAM constructor in block.py. Let me know if further clarification is needed!

[AnirudhSanthana]

Sure!! I will look into this further and revert back after debugging....it’s possible that there something we are missing. Meanwhile, can you please confirm whether the results shown in this image are from the validation set or the training set?

As in the YAML file, we typically provide paths for all three datasets: training, validation, and testing. However, I couldn’t find the results for the test dataset, even in the output folder. Could you clarify where those results are stored or how to access them?

Would appreciate your guidance on this! @glenn-jocher

[glenn-jocher]

@AnirudhSanthana hi Anirudh, For test set results, you'll need to explicitly run validation on your test split using model.val(data='your_data.yaml', split='test'). Test results aren't generated automatically during training - only validation metrics are shown by default. The output will be saved in your runs/detect/val directory. For more details on dataset splits and validation, see our documentation at https://docs.ultralytics.com/modes/val/.