Transformer in Transformer 开源项目教程

最新推荐文章于 2025-04-14 17:00:59 发布

原创最新推荐文章于 2025-04-14 17:00:59 发布 · 551 阅读

17 ·

CC 4.0 BY-SA版权

Transformer in Transformer 开源项目教程

项目地址:https://gitcode.com/gh_mirrors/tr/transformer-in-transformer

项目介绍

Transformer in Transformer (TNT) 是一种新型的神经网络架构，它通过结合像素级和块级注意力机制来增强图像分类的性能。TNT 项目由 Kai Han 等人提出，并在 GitHub 上开源，由 lucidrains 实现。该项目的主要创新在于将图像分割成更小的块（称为“视觉单词”），并计算这些小块与其他小块之间的注意力，从而提高特征提取的精细度。

项目快速启动

安装

首先，确保你已经安装了 PyTorch。然后，通过 pip 安装 TNT 模型：

pip install transformer-in-transformer

使用示例

以下是一个简单的代码示例，展示如何使用 TNT 模型进行图像分类：

import torch
from transformer_in_transformer import TNT

# 初始化 TNT 模型
tnt = TNT(
    image_size=256,  # 图像大小
    patch_dim=512,   # 块维度
    pixel_dim=24,    # 像素维度
    patch_size=16,   # 块大小
    pixel_size=4,    # 像素大小
    depth=6,         # 深度
    num_classes=1000,  # 输出类别数
    attn_dropout=0.1,  # 注意力 dropout
    ff_dropout=0.1     # 前馈 dropout
)

# 生成随机图像数据
img = torch.randn(2, 3, 256, 256)

# 获取分类结果
logits = tnt(img)  # 输出形状为 (2, 1000)
print(logits)