用于图像分类的全局滤波网络

由 Yongming Rao、Wenliang Zhao、Zheng Zhu、Jiwen Lu、Jie Zhou 创建

本仓库包含 GFNet（NeurIPS 2021 & T-PAMI）的 PyTorch 实现。

全局滤波网络是一种基于 Transformer 的架构，它以对数线性复杂度在频域中学习长期的空间依赖关系。我们的架构用三个关键操作替换了视觉 Transformer 中的自注意力层：二维离散傅里叶变换、频域特征与可学习全局滤波器之间的逐元素乘法，以及二维逆傅里叶变换。

我们的代码基于 pytorch-image-models 和 DeiT。

[项目页面] [arXiv]

全局滤波层

GFNet 是一种概念简单但计算高效的架构，由多个堆叠的全局滤波层和前馈网络（FFN）组成。全局滤波层利用高效快速傅里叶变换（FFT）算法，以对数线性的复杂度混合特征图中的 token。该层实现起来非常容易：

import torch
import torch.nn as nn
import torch.fft

class GlobalFilter(nn.Module):
    def __init__(self, dim, h=14, w=8):
        super().__init__()
        self.complex_weight = nn.Parameter(torch.randn(h, w, dim, 2, dtype=torch.float32) * 0.02)

    def forward(self, x):
        B, H, W, C = x.shape
        x = torch.fft.rfft2(x, dim=(1, 2), norm='ortho')
        weight = torch.view_as_complex(self.complex_weight)
        x = x * weight
        x = torch.fft.irfft2(x, s=(H, W), dim=(1, 2), norm='ortho')
        return x

与自注意力机制和空间 MLP 相比，我们的全局滤波层在处理高分辨率特征图时效率更高：

模型库

我们提供了在 ImageNet 上预训练的 GFNet 模型：

名称	架构	参数量	FLOPs	top-1 精确度	top-5 精确度	下载链接
GFNet-Ti	gfnet-ti	7M	1.3G	74.6	92.2	清华大学云 / Google Drive
GFNet-XS	gfnet-xs	16M	2.8G	78.6	94.2	清华大学云 / Google Drive
GFNet-S	gfnet-s	25M	4.5G	80.0	94.9	清华大学云 / Google Drive
GFNet-B	gfnet-b	43M	7.9G	80.7	95.1	清华大学云 / Google Drive
GFNet-H-Ti	gfnet-h-ti	15M	2.0G	80.1	95.1	清华大学云 / Google Drive
GFNet-H-S	gfnet-h-s	32M	4.5G	81.5	95.6	清华大学云 / Google Drive
GFNet-H-B	gfnet-h-b	54M	8.4G	82.9	96.2	清华大学云 / Google Drive

使用方法

需求

• torch>=1.8.0
• torchvision
• timm

注: 要使用 PyTorch 中的 rfft2 和 irfft2 函数，需要安装 PyTorch>=1.8.0。虽然 PyTorch 1.6.0 之后就支持复数运算，但 fft API 在不同版本间存在细微差异。

数据准备: 从 http://image-net.org/ 下载并解压 ImageNet 图片。目录结构应如下所示：

│ILSVRC2012/
├──train/
│  ├── n01440764
│  │   ├── n01440764_10026.JPEG
│  │   ├── n01440764_10027.JPEG
│  │   ├── ......
│  ├── ......
├──val/
│  ├── n01440764
│  │   ├── ILSVRC2012_val_00000293.JPEG
│  │   ├── ILSVRC2012_val_00002138.JPEG
│  │   ├── ......
│  ├── ......

评估

要在单个 GPU 上评估预训练的 GFNet 模型在 ImageNet 验证集上的性能，请运行：

python infer.py --data-path /path/to/ILSVRC2012/ --arch arch_name --model-path /path/to/model

训练

ImageNet

要从头开始在 ImageNet 上训练 GFNet 模型，请运行：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main_gfnet.py  --output_dir logs/gfnet-xs --arch gfnet-xs --batch-size 128 --data-path /path/to/ILSVRC2012/

若需在更高分辨率下微调预训练模型，请运行：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main_gfnet.py  --output_dir logs/gfnet-xs-img384 --arch gfnet-xs --input-size 384 --batch-size 64 --data-path /path/to/ILSVRC2012/ --lr 5e-6 --weight-decay 1e-8 --min-lr 5e-6 --epochs 30 --finetune /path/to/model

迁移学习数据集

要在迁移学习数据集上微调预训练模型，请运行：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main_gfnet_transfer.py  --output_dir logs/gfnet-xs-cars --arch gfnet-xs --batch-size 64 --data-set CARS --data-path /path/to/stanford_cars --epochs 1000 --lr 0.0001 --weight-decay 1e-4 --clip-grad 1 --warmup-epochs 5 --finetune /path/to/model 

可视化

为了更直观地理解我们的全局滤波操作，我们可视化了 GFNet-XS 不同层中学习到的滤波器。

许可证

MIT 许可证

引用

如果您在研究中使用了我们的工作，请考虑引用以下文献：

@inproceedings{rao2021global,
  title={Global Filter Networks for Image Classification},
  author={Rao, Yongming and Zhao, Wenliang and Zhu, Zheng and Lu, Jiwen and Zhou, Jie},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year = {2021}
}

GFNet 快速上手指南

GFNet (Global Filter Networks) 是一种基于频域操作的 Transformer 风格架构，通过二维离散傅里叶变换（2D DFT）和可学习的全局滤波器来捕捉长程空间依赖，具有对数线性复杂度，在处理高分辨率图像时效率显著。

环境准备

系统要求

• Python: 建议 3.7+
• PyTorch: >= 1.8.0 (必须，以支持 rfft2 和 irfft2 API)
• GPU: 推荐用于训练和推理（支持 CUDA）

前置依赖

请确保安装以下核心库：

• torch
• torchvision
• timm (pytorch-image-models)

注意：虽然 PyTorch 1.6.0+ 支持复数运算，但本项目使用的 FFT API 接口在 1.8.0 版本中才稳定，请务必升级。

安装步骤

1. 克隆项目代码

   git clone https://github.com/raoyongming/GFNet.git
   cd GFNet
   

2. 安装依赖 建议使用国内镜像源加速安装（如清华源）：

   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
   pip install timm -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   (注：cu118 请根据你的实际 CUDA 版本调整，若无 GPU 可移除 --index-url 部分或使用 cpu 版本)

3. 准备数据集 (ImageNet) 下载 ImageNet 数据集并解压，目录结构需如下所示：

   ILSVRC2012/
   ├── train/
   │   ├── n01440764/
   │   │   ├── n01440764_10026.JPEG
   │   │   └── ...
   ├── val/
   │   ├── n01440764/
   │   │   ├── ILSVRC2012_val_00000293.JPEG
   │   │   └── ...
   

基本使用

1. 模型推理 (Evaluation)

使用预训练模型在 ImageNet 验证集上进行单卡推理测试。

python infer.py --data-path /path/to/ILSVRC2012/ --arch gfnet-xs --model-path /path/to/model_checkpoint.pth

请将 /path/to/... 替换为实际的数据集路径和下载的模型权重路径。

2. 从头训练 (Training)

在 ImageNet 上从头训练 GFNet-XS 模型（默认使用 8 张 GPU 分布式训练）：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main_gfnet.py  --output_dir logs/gfnet-xs --arch gfnet-xs --batch-size 128 --data-path /path/to/ILSVRC2012/

3. 微调高分辨率模型 (Finetune)

加载预训练权重，在更高分辨率（如 384x384）下进行微调：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main_gfnet.py  --output_dir logs/gfnet-xs-img384 --arch gfnet-xs --input-size 384 --batch-size 64 --data-path /path/to/ILSVRC2012/ --lr 5e-6 --weight-decay 1e-8 --min-lr 5e-6 --epochs 30 --finetune /path/to/pretrained_model.pth

4. 迁移学习 (Transfer Learning)

在其他数据集（如 Stanford Cars）上进行微调：

python -m torch.distributed.launch --nproc_per_node=8 --use_env main_gfnet_transfer.py  --output_dir logs/gfnet-xs-cars --arch gfnet-xs --batch-size 64 --data-set CARS --data-path /path/to/stanford_cars --epochs 1000 --lr 0.0001 --weight-decay 1e-4 --clip-grad 1 --warmup-epochs 5 --finetune /path/to/pretrained_model.pth

---

预训练模型下载： 项目提供多种规格模型（Ti, XS, S, B 等），推荐优先使用清华云盘链接下载以获得更快速度。具体链接请参考原项目 README 中的 "Model Zoo" 表格。