Ultralytics：解读TransformerLayer模块

原创 于 2026-06-30 00:00:00 发布 · 242 阅读 · 3 ·
本内容遵循CC 4.0 BY-SA版权协议
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
GEO检测
标签
#python #pytorch #YOLO #人工智能 #深度学习

Ultralytics：解读TransformerLayer模块

前言

相关介绍

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

TransformerLayer（简化版Transformer层）

TransformerLayer 是一个极简的Transformer编码器层，去除了LayerNorm和Dropout，仅保留核心的 多头自注意力（Multi-Head Attention） 和 前馈网络（FFN），并采用残差连接。它参考了论文 “An Image is Worth 16x16 Words” 中的设计，在Ultralytics中用于轻量级的特征变换，例如YOLOv8中的某些注意力模块。由于没有LayerNorm，该层在训练时可能对初始化更敏感，但也减少了计算开销。

代码实现

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

class TransformerLayer(nn.Module):
    """Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)."""

    def __init__(self, c: int, num_heads: int):
        """Initialize a self-attention mechanism using linear transformations and multi-head attention.

        Args:
            c (int): Input and output channel dimension.
            num_heads (int): Number of attention heads.
        """
        super().__init__()
        self.q = nn.Linear(c, c, bias=False)
        self.k = nn.Linear(c, c, bias=False)
        self.v = nn.Linear(c, c, bias=False)
        self.ma = nn.MultiheadAttention(embed_dim=c, num_heads=num_heads)
        self.fc1 = nn.Linear(c, c, bias=False)
        self.fc2 = nn.Linear(c, c, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply a transformer block to the input x and return the output.

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

        Returns:
            (torch.Tensor): Output tensor after transformer layer.
        """
        x = self.ma(self.q(x), self.k(x), self.v(x))[0] + x
        return self.fc2(self.fc1(x)) + x

功能

  • 自注意力特征变换：通过多头自注意力捕捉序列中token之间的长距离依赖关系，并通过前馈网络进行非线性变换。
  • 残差连接：注意力输出和FFN输出均与输入相加，缓解梯度消失并加速训练。
  • 无归一化层：去除了LayerNorm，旨在提升推理速度（根据注释），但训练时需注意梯度稳定性。
  • 无偏置：所有线性层均不带偏置，进一步减少参数量。

初始化参数

参数类型说明
cint输入/输出特征维度（必须与 embed_dim 一致）
num_headsint多头注意力的头数（c 必须能被 num_heads 整除）

前向方法

  • forward(x)：输入 x，输出与 x 形状相同的张量。

细节：

  1. 对 x 分别计算查询（Q）、键（K）、值（V）投影。
  2. 送入 nn.MultiheadAttention（未指定 batch_first，因此输入形状需为 (seq_len, batch, c)）。
  3. 注意力输出与原始 x 相加（残差1）。
  4. 经过两个全连接层（fc1 → fc2，中间无激活函数？实际上代码中没有激活函数，仅有线性变换，所以是 fc2(fc1(x))，这相当于一个两层的线性变换，但仍包含非线性能力？实际上没有激活函数，所以等价于一个线性层，但这里可能笔误？通常FFN包括激活函数。但这里没有，所以可能是一个简化，也可以认为 fc1 和 fc2 共同构成一个线性变换，但两层线性相当于一层，除非中间有激活。没有激活，所以这个FFN实际上退化为线性层。但设计者可能意图是用 fc1 做升维、fc2 降维，但未加激活，所以整体仍是线性的。然而，注释说“LayerNorm layers removed for better performance”，但未提及去除激活，可能是疏忽？但代码确实没有激活。所以我们按代码解释。）
  5. 残差连接：x = fc2(fc1(x)) + x。

注意：由于没有激活函数，FFN部分实际上是线性变换，但两次线性变换可能等效于一次线性变换（如果中间维度相同），但这里两个线性层都是 c 到 c，无激活，所以 fc2(fc1(x)) 相当于一个线性变换（复合线性），但乘以两个矩阵，可能是为了增加参数量？但无激活并不改变线性性质。

使用示例

if __name__ == '__main__':
    # 设置参数
    seq_len, batch, dim, heads = 10, 2, 128, 4
    # 创建随机输入（按默认的 (seq_len, batch, dim) 格式）
    x = torch.randn(seq_len, batch, dim)

    # 创建 TransformerLayer
    layer = TransformerLayer(c=dim, num_heads=heads)

    # 前向传播
    with torch.no_grad():
        out = layer(x)
    print("输入形状:", x.shape)   # [10, 2, 128]
    print("输出形状:", out.shape) # [10, 2, 128]

    # 若使用 batch_first=True，需手动调整或修改
    # 但这里未设，所以默认。

输出示例：

输入形状: torch.Size([10, 2, 128])
输出形状: torch.Size([10, 2, 128])

流程示意图

在这里插入图片描述

代码解读

__init__ 方法
  • self.q, self.k, self.v：三个独立的线性层，无偏置，用于将输入映射为查询、键、值。
  • self.ma：nn.MultiheadAttention，默认 batch_first=False，因此输入形状需为 (seq_len, batch, embed_dim)。
  • self.fc1, self.fc2：两个线性层，无偏置，无激活函数，构成前馈部分（通常应包含激活，但这里被省略）。
forward 方法
  1. 计算 q(x), k(x), v(x)，送入 self.ma，取第一个返回值（注意力输出），忽略注意力权重。
  2. 残差相加：x = x + attn_output。
  3. 通过 fc1 和 fc2（无激活），再次残差相加。
  4. 返回结果。

注意事项

  1. 输入格式：由于 nn.MultiheadAttention 默认 batch_first=False，输入应为 (seq_len, batch, c)。若需 batch_first=True，需在初始化中设置或自行调整。
  2. 无LayerNorm：该层不包含归一化，训练时可能出现梯度不稳定，需谨慎设置学习率，或使用预训练权重初始化。
  3. 无Dropout：过拟合风险可能增加，尤其在数据量较小时。
  4. 无激活函数：fc1 和 fc2 之间没有激活函数，FFN实际为线性变换，限制了模型的非线性表达能力。这可能是设计上的简化，但降低了模型容量。
  5. 位置编码：该模块本身不包含位置信息，若用于序列建模（如视觉特征），需在输入前添加位置编码。
  6. 参数量：所有线性层无偏置，参数量较少，适合轻量级部署。

优缺点

优点
  1. 计算高效：去除LayerNorm、Dropout和激活函数，减少了前向/反向传播的计算开销。
  2. 参数量少：无偏置，模型更轻量，适合移动端或边缘设备。
  3. 实现简洁：代码量极少，易于理解和调试。
  4. 残差结构：有助于梯度传播，支持更深的堆叠。
缺点
  1. 表达能力受限：缺乏LayerNorm可能导致训练不稳定，缺乏激活函数使得FFN退化为线性变换，削弱了模型的非线性拟合能力。
  2. 缺少正则化：无Dropout，容易过拟合，尤其在小数据集上。
  3. 位置缺失：不包含位置编码，需要外部提供，否则无法利用序列顺序信息。
  4. 通用性不足：该简化设计可能仅适用于特定任务（如特定尺寸的视觉特征），迁移到其他场景时需额外修改。

在YOLOv8等模型中，TransformerLayer 常用于构造更复杂的注意力模块（如 TransformerBlock）。实际使用时，建议根据任务需求考虑是否添加LayerNorm和激活函数，以平衡性能与效率。

参考文献

[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、付费专栏及课程。

余额充值