Ultralytics:解读C3k2模块

在这里插入图片描述

前言

相关介绍

Ultralytics 简介

Ultralytics 基于多年的计算机视觉和人工智能基础研究,创建了最先进的 (SOTA) YOLO 模型。我们的模型不断更新性能和灵活性,快速、准确且易于使用。他们擅长对象检测、跟踪、实例分割、语义分割、图像分类和姿势估计任务。

前提条件

  • 熟悉Python、Pytorch

实验环境

Package                  Version
------------------------ ------------
Python                   3.11.8
absl-py                  2.4.0
accelerate               1.13.0
annotated-doc            0.0.4
anyio                    4.13.0
calflops                 0.3.2
certifi                  2026.4.22
charset-normalizer       3.4.7
click                    8.3.3
colorama                 0.4.6
contourpy                1.3.3
cycler                   0.12.1
filelock                 3.29.0
flatbuffers              25.12.19
fonttools                4.62.1
fsspec                   2026.4.0
grpcio                   1.80.0
h11                      0.16.0
hf-xet                   1.5.0
httpcore                 1.0.9
httpx                    0.28.1
huggingface_hub          1.14.0
idna                     3.15
Jinja2                   3.1.6
kiwisolver               1.5.0
Markdown                 3.10.2
markdown-it-py           4.2.0
MarkupSafe               3.0.3
matplotlib               3.10.9
mdurl                    0.1.2
ml_dtypes                0.5.0
mpmath                   1.3.0
networkx                 3.6.1
numpy                    1.26.4
nvidia-cublas-cu12       12.8.3.14
nvidia-cuda-cupti-cu12   12.8.57
nvidia-cuda-nvrtc-cu12   12.8.61
nvidia-cuda-runtime-cu12 12.8.57
nvidia-cudnn-cu12        9.7.1.26
nvidia-cufft-cu12        11.3.3.41
nvidia-cufile-cu12       1.13.0.11
nvidia-curand-cu12       10.3.9.55
nvidia-cusolver-cu12     11.7.2.55
nvidia-cusparse-cu12     12.5.7.53
nvidia-cusparselt-cu12   0.6.3
nvidia-nccl-cu12         2.26.2
nvidia-nvjitlink-cu12    12.8.61
nvidia-nvtx-cu12         12.8.55
onnx                     1.19.0
onnxruntime-gpu          1.26.0
onnxslim                 0.1.94
opencv-python            4.6.0.66
packaging                26.2
pillow                   12.2.0
pip                      24.0
polars                   1.40.1
polars-runtime-32        1.40.1
protobuf                 7.34.1
psutil                   7.2.2
pycocotools              2.0.11
Pygments                 2.20.0
pyparsing                3.3.2
python-dateutil          2.9.0.post0
PyYAML                   6.0.3
regex                    2026.5.9
requests                 2.34.1
rich                     15.0.0
safetensors              0.7.0
scipy                    1.16.0
setuptools               65.5.0
shellingham              1.5.4
six                      1.17.0
sympy                    1.14.0
tabulate                 0.10.0
tensorboard              2.20.0
tensorboard-data-server  0.7.2
tokenizers               0.22.2
torch                    2.7.1+cu128
torchaudio               2.7.1+cu128
torchvision              0.22.1+cu128
tqdm                     4.67.3
transformers             5.8.1
triton                   3.3.1
typer                    0.25.1
typing_extensions        4.15.0
ultralytics              8.4.58
ultralytics-thop         2.0.19
urllib3                  2.7.0
Werkzeug                 3.1.8

C3k2(灵活的 CSP 模块,支持多种子块组合)

C3k2C2f 的增强版本,它允许用户在每个块中灵活选择不同的子模块组合,包括 标准 BottleneckC3k(可定制核大小的 CSP 块)Bottleneck + PSABlock(自注意力增强)。这种设计使得网络能够根据任务需求在精度和速度之间进行更精细的权衡,同时保持了 C2f 的高效级联结构。


代码实现

import cv2
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt

def autopad(k, p=None, d=1):  # kernel, padding, dilation
    """Pad to 'same' shape outputs."""
    if d > 1:
        k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]  # actual kernel-size
    if p is None:
        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
    return p

class Conv(nn.Module):
    """Standard convolution module with batch normalization and activation.

    Attributes:
        conv (nn.Conv2d): Convolutional layer.
        bn (nn.BatchNorm2d): Batch normalization layer.
        act (nn.Module): Activation function layer.
        default_act (nn.Module): Default activation function (SiLU).
    """

    default_act = nn.SiLU()  # default activation

    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
        """Initialize Conv layer with given parameters.

        Args:
            c1 (int): Number of input channels.
            c2 (int): Number of output channels.
            k (int): Kernel size.
            s (int): Stride.
            p (int, optional): Padding.
            g (int): Groups.
            d (int): Dilation.
            act (bool | nn.Module): Activation function.
        """
        super().__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()

    def forward(self, x):
        """Apply convolution, batch normalization and activation to input tensor.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            (torch.Tensor): Output tensor.
        """
        return self.act(self.bn(self.conv(x)))

    def forward_fuse(self, x):
        """Apply convolution and activation without batch normalization.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            (torch.Tensor): Output tensor.
        """
        return self.act(self.conv(x))

class Bottleneck(nn.Module):
    """Standard bottleneck."""

    def __init__(
        self, c1: int, c2: int, shortcut: bool = True, g: int = 1, k: tuple[int, int] = (3, 3), e: float = 0.5
    ):
        """Initialize a standard bottleneck module.

        Args:
            c1 (int): Input channels.
            c2 (int): Output channels.
            shortcut (bool): Whether to use shortcut connection.
            g (int): Groups for convolutions.
            k (tuple): Kernel sizes for convolutions.
            e (float): Expansion ratio.
        """
        super().__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, c_, k[0], 1)
        self.cv2 = Conv(c_, c2, k[1], 1, g=g)
        self.add = shortcut and c1 == c2

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply bottleneck with optional shortcut connection."""
        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))

class C2f(nn.Module):
    """Faster Implementation of CSP Bottleneck with 2 convolutions."""

    def __init__(self, c1: int, c2: int, n: int = 1, shortcut: bool = False, g: int = 1, e: float = 0.5):
        """Initialize a CSP bottleneck with 2 convolutions.

        Args:
            c1 (int): Input channels.
            c2 (int): Output channels.
            n (int): Number of Bottleneck blocks.
            shortcut (bool): Whether to use shortcut connections.
            g (int): Groups for convolutions.
            e (float): Expansion ratio.
        """
        super().__init__()
        self.c = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, 2 * self.c, 1, 1)
        self.cv2 = Conv((2 + n) * self.c, c2, 1)  # optional act=FReLU(c2)
        self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through C2f layer."""
        y = list(self.cv1(x).chunk(2, 1))
        y.extend(m(y[-1]) for m in self.m)
        return self.cv2(torch.cat(y, 1))

    def forward_split(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass using split() instead of chunk()."""
        y = self.cv1(x).split((self.c, self.c), 1)
        y = [y[0], y[1]]
        y.extend(m(y[-1]) for m in self.m)
        return self.cv2(torch.cat(y, 1))
    
class C3(nn.Module):
    """CSP Bottleneck with 3 convolutions."""

    def __init__(self, c1: int, c2: int, n: int = 1, shortcut: bool = True, g: int = 1, e: float = 0.5):
        """Initialize the CSP Bottleneck with 3 convolutions.

        Args:
            c1 (int): Input channels.
            c2 (int): Output channels.
            n (int): Number of Bottleneck blocks.
            shortcut (bool): Whether to use shortcut connections.
            g (int): Groups for convolutions.
            e (float): Expansion ratio.
        """
        super().__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c1, c_, 1, 1)
        self.cv3 = Conv(2 * c_, c2, 1)  # optional act=FReLU(c2)
        self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, k=((1, 1), (3, 3)), e=1.0) for _ in range(n)))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through the CSP bottleneck with 3 convolutions."""
        return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))

class C3k(C3):
    """C3k is a CSP bottleneck module with customizable kernel sizes for feature extraction in neural networks."""

    def __init__(self, c1: int, c2: int, n: int = 1, shortcut: bool = True, g: int = 1, e: float = 0.5, k: int = 3):
        """Initialize C3k module.

        Args:
            c1 (int): Input channels.
            c2 (int): Output channels.
            n (int): Number of Bottleneck blocks.
            shortcut (bool): Whether to use shortcut connections.
            g (int): Groups for convolutions.
            e (float): Expansion ratio.
            k (int): Kernel size.
        """
        super().__init__(c1, c2, n, shortcut, g, e)
        c_ = int(c2 * e)  # hidden channels
        # self.m = nn.Sequential(*(RepBottleneck(c_, c_, shortcut, g, k=(k, k), e=1.0) for _ in range(n)))
        self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, k=(k, k), e=1.0) for _ in range(n)))

class Attention(nn.Module):
    """Attention module that performs self-attention on the input tensor.

    Args:
        dim (int): The input tensor dimension.
        num_heads (int): The number of attention heads.
        attn_ratio (float): The ratio of the attention key dimension to the head dimension.

    Attributes:
        num_heads (int): The number of attention heads.
        head_dim (int): The dimension of each attention head.
        key_dim (int): The dimension of the attention key.
        scale (float): The scaling factor for the attention scores.
        qkv (Conv): Convolutional layer for computing the query, key, and value.
        proj (Conv): Convolutional layer for projecting the attended values.
        pe (Conv): Convolutional layer for positional encoding.
    """

    def __init__(self, dim: int, num_heads: int = 8, attn_ratio: float = 0.5):
        """Initialize multi-head attention module.

        Args:
            dim (int): Input dimension.
            num_heads (int): Number of attention heads.
            attn_ratio (float): Attention ratio for key dimension.
        """
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.key_dim = int(self.head_dim * attn_ratio)
        self.scale = self.key_dim**-0.5
        nh_kd = self.key_dim * num_heads
        h = dim + nh_kd * 2
        self.qkv = Conv(dim, h, 1, act=False)
        self.proj = Conv(dim, dim, 1, act=False)
        self.pe = Conv(dim, dim, 3, 1, g=dim, act=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass of the Attention module.

        Args:
            x (torch.Tensor): The input tensor.

        Returns:
            (torch.Tensor): The output tensor after self-attention.
        """
        B, C, H, W = x.shape
        N = H * W
        qkv = self.qkv(x)
        q, k, v = qkv.view(B, self.num_heads, self.key_dim * 2 + self.head_dim, N).split(
            [self.key_dim, self.key_dim, self.head_dim], dim=2
        )

        attn = (q.transpose(-2, -1) @ k) * self.scale
        attn = attn.softmax(dim=-1)
        x = (v @ attn.transpose(-2, -1)).view(B, C, H, W) + self.pe(v.reshape(B, C, H, W))
        x = self.proj(x)
        return x
    
class PSABlock(nn.Module):
    """PSABlock class implementing a Position-Sensitive Attention block for neural networks.

    This class encapsulates the functionality for applying multi-head attention and feed-forward neural network layers
    with optional shortcut connections.

    Attributes:
        attn (Attention): Multi-head attention module.
        ffn (nn.Sequential): Feed-forward neural network module.
        add (bool): Flag indicating whether to add shortcut connections.

    Methods:
        forward: Performs a forward pass through the PSABlock, applying attention and feed-forward layers.

    Examples:
        Create a PSABlock and perform a forward pass
        >>> psablock = PSABlock(c=128, attn_ratio=0.5, num_heads=4, shortcut=True)
        >>> input_tensor = torch.randn(1, 128, 32, 32)
        >>> output_tensor = psablock(input_tensor)
    """

    def __init__(self, c: int, attn_ratio: float = 0.5, num_heads: int = 4, shortcut: bool = True) -> None:
        """Initialize the PSABlock.

        Args:
            c (int): Input and output channels.
            attn_ratio (float): Attention ratio for key dimension.
            num_heads (int): Number of attention heads.
            shortcut (bool): Whether to use shortcut connections.
        """
        super().__init__()

        self.attn = Attention(c, attn_ratio=attn_ratio, num_heads=num_heads)
        self.ffn = nn.Sequential(Conv(c, c * 2, 1), Conv(c * 2, c, 1, act=False))
        self.add = shortcut

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Execute a forward pass through PSABlock.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            (torch.Tensor): Output tensor after attention and feed-forward processing.
        """
        x = x + self.attn(x) if self.add else self.attn(x)
        x = x + self.ffn(x) if self.add else self.ffn(x)
        return x

class C3k2(C2f):
    """Faster Implementation of CSP Bottleneck with 2 convolutions."""

    def __init__(
        self,
        c1: int,
        c2: int,
        n: int = 1,
        c3k: bool = False,
        e: float = 0.5,
        attn: bool = False,
        g: int = 1,
        shortcut: bool = True,
    ):
        """Initialize C3k2 module.

        Args:
            c1 (int): Input channels.
            c2 (int): Output channels.
            n (int): Number of blocks.
            c3k (bool): Whether to use C3k blocks.
            e (float): Expansion ratio.
            attn (bool): Whether to use attention blocks.
            g (int): Groups for convolutions.
            shortcut (bool): Whether to use shortcut connections.
        """
        super().__init__(c1, c2, n, shortcut, g, e)
        self.m = nn.ModuleList(
            nn.Sequential(
                Bottleneck(self.c, self.c, shortcut, g),
                PSABlock(self.c, attn_ratio=0.5, num_heads=max(self.c // 64, 1)),
            )
            if attn
            else C3k(self.c, self.c, 2, shortcut, g)
            if c3k
            else Bottleneck(self.c, self.c, shortcut, g)
            for _ in range(n)
        )

功能

  • CSP 级联结构:继承自 C2f,通过 cv1 将输入通道扩展为 2*self.c,拆分后一路直接保留,另一路依次通过 self.m 中的多个块(每个块类型可单独选择),每次输出追加到特征列表,最终拼接并通过 cv2 压缩到目标通道。
  • 可配置块类型
    • attn=True:每个块由 Bottleneck + PSABlock 串联组成,先通过标准瓶颈提取特征,再经自注意力增强全局上下文。
    • c3k=True:每个块为 C3k 模块(内部包含两个 Bottleneck,核大小默认为 3,隐藏通道减半),提供更深的特征变换。
    • 默认(两者均为 False):每个块为单个 Bottleneck,与 C2f 相同。
  • 灵活组合:所有块共享相同类型,不能混合不同种类,但可通过调整 n 控制深度。

初始化参数

参数类型说明
c1int输入通道数
c2int输出通道数
nint块的数量(默认 1)
c3kbool是否使用 C3k 作为块(默认 False)
efloat扩展比,self.c = int(c2 * e)(默认 0.5)
attnbool是否使用 Bottleneck + PSABlock 作为块(默认 False)
gint分组卷积组数(作用于 BottleneckC3k 内部的卷积)
shortcutbool是否在 BottleneckC3k 内部使用残差连接(默认 True)
  • self.c 由父类 C2f 计算得到。
  • attn=Truec3k=True,由于 if attn 优先,会使用注意力组合,忽略 c3k
  • C3k 在调用时使用默认 k=3(核大小),若需要自定义核,可修改源码或继承重写。

前向方法

  • forward(x):继承自 C2f,流程与父类一致。
  • forward_split(x):继承自 C2f,使用 split 替代 chunk

数据流(与 C2f 相同):

  1. y = self.cv1(x)[B, 2*self.c, H, W]
  2. 拆分为 [y0, y1],各 [B, self.c, H, W]
  3. self.m 中的每个块,取当前列表最后一个元素输入,输出追加到列表
  4. 拼接所有 (2+n) 个张量
  5. 通过 self.cv2 压缩到 c2

使用示例

在这里插入图片描述

if __name__ == '__main__':
    # 1. 随机输入
    B, c1, c2, H, W = 2, 32, 64, 64, 64
    x = torch.randn(B, c1, H, W)

    # 2. 不同配置的 C3k2
    # 配置1:默认(仅 Bottleneck)
    model1 = C3k2(c1, c2, n=2, c3k=False, attn=False, e=0.5, shortcut=True)
    with torch.no_grad():
        out1 = model1(x)
    print("默认输出形状:", out1.shape)  # [2, 64, 64, 64]

    # 配置2:使用 C3k 块
    model2 = C3k2(c1, c2, n=2, c3k=True, attn=False, e=0.5)
    with torch.no_grad():
        out2 = model2(x)
    print("C3k 输出形状:", out2.shape)  # [2, 64, 64, 64]

    # 配置3:使用 Bottleneck + PSABlock
    model3 = C3k2(c1, c2, n=2, c3k=False, attn=True, e=0.5)
    with torch.no_grad():
        out3 = model3(x)
    print("注意力输出形状:", out3.shape)  # [2, 64, 64, 64]

    # 4. 真实图像演示(单通道扩展)
    img_path = "cat_640x640.png"
    img_bgr = cv2.imread(img_path)
    if img_bgr is not None:
        img_gray = cv2.cvtColor(cv2.resize(img_bgr, (64, 64)), cv2.COLOR_BGR2GRAY)
        img_tensor = torch.from_numpy(img_gray).float().unsqueeze(0).unsqueeze(0)  # [1,1,64,64]
        # 扩展至 c1=32
        x_img = img_tensor.repeat(1, c1, 1, 1)

        # 使用配置2(C3k)演示
        model_demo = C3k2(c1=32, c2=32, n=1, c3k=True, e=0.5)
        with torch.no_grad():
            out_img = model_demo(x_img)

        inp_ch0 = x_img[0, 0].cpu().numpy()
        out_ch0 = out_img[0, 0].cpu().numpy()

        def norm(arr):
            return (arr - arr.min()) / (arr.max() - arr.min() + 1e-8)

        plt.figure(figsize=(12, 5), constrained_layout=True)
        plt.subplot(1, 3, 1)
        plt.imshow(img_gray, cmap='gray')
        plt.title("Original")
        plt.axis("off")
        plt.subplot(1, 3, 2)
        plt.imshow(norm(inp_ch0), cmap='gray')
        plt.title("Input Ch0")
        plt.axis("off")
        plt.subplot(1, 3, 3)
        plt.imshow(norm(out_ch0), cmap='gray')
        plt.title("C3k2 (C3k) Output")
        plt.axis("off")
        plt.savefig("c3k2_demo.png", dpi=150)
        print("可视化已保存为 c3k2_demo.png")

在这里插入图片描述
输出示例

默认输出形状: torch.Size([2, 64, 64, 64])
C3k 输出形状: torch.Size([2, 64, 64, 64])
注意力输出形状: torch.Size([2, 64, 64, 64])
可视化已保存为 c3k2_demo.png

流程示意图

输入 x (B, c1, H, W)

Conv 1x1: c1 → 2*c (cv1)

chunk(2,1): y0, y1

y0 直接保留

y1

块 1
(类型由参数决定)

输出 y2

块 2

...

块 n

输出 y_{n+1}

拼接: y0, y1, ..., y_{n+1}

Conv 1x1: (2+n)*c → c2 (cv2)

输出 (B, c2, H, W)


代码解读

  • __init__
    • 调用父类 C2f.__init__(c1, c2, n, shortcut, g, e),创建 self.cself.cv1self.cv2 以及父类的 self.m
    • 根据参数 attnc3k 构造新的 self.m
      • attn=True:每个块为 nn.Sequential(Bottleneck(...), PSABlock(...))
      • 否则若 c3k=True:每个块为 C3k(self.c, self.c, 2, shortcut, g)(注意 n=2,即每个 C3k 包含两个 Bottleneck)。
      • 否则:每个块为单个 Bottleneck(self.c, self.c, shortcut, g)
    • 列表推导生成 n 个块。
  • forward / forward_split:继承自 C2f,无需重写。

注意事项

  1. attnc3k 互斥attn 优先级更高,若同时为 True,则忽略 c3k
  2. C3k 的默认参数:在 c3k=True 时,C3k 使用默认 e=0.5,会导致其内部隐藏通道减半(c_ = c // 2),可能改变计算量。若需调整,可修改源码。
  3. PSABlock 的头数num_heads = max(self.c // 64, 1),当 self.c < 64 时头数为 1,可能影响注意力效果。
  4. 通道约束self.c = int(c2 * e) 必须为正整数,且 C3k 内部要求 c_ 为整数。
  5. 训练与推理差异:若使用 PSABlock,其内部包含自注意力,训练时需注意学习率和正则化;推理时无特殊操作。
  6. C2f 的兼容性:接口与 C2f 相似,但增加了 attnc3k 参数,可灵活替换。

优缺点

优点
  1. 高度灵活:通过简单开关选择不同块类型,适应不同精度/速度需求。
  2. 组合丰富:可将标准卷积、CSP 扩展和自注意力混合使用,提升表达能力。
  3. 继承 C2f 高效结构:级联梯度流和轻量设计依然保留。
  4. 即插即用:与 C2f 完全兼容,只需传递额外参数。
缺点
  1. 参数组合复杂attnc3k 的优先级需明确,容易误用。
  2. C3k 的默认设置可能不最优e=0.5k=3 固定,若需自定义需修改代码。
  3. 注意力块增加计算开销PSABlock 包含自注意力,在高分辨率时显存占用大。
  4. 所有块类型相同:不能在同一个 C3k2 中混合不同块类型,降低了组合灵活性。

在 YOLOv8 系列中,C3k2 可用于替换 C2f 以提供更多配置选项。建议根据任务需求选择 attnc3k,并在小分辨率特征图(如 16×16)上启用注意力,以获得最佳平衡。

参考文献

[1] https://docs.ultralytics.com/
[2] https://github.com/ultralytics/ultralytics.git

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

FriendshipT

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值