焦点频率损失 - 官方 PyTorch 实现

本仓库提供了以下论文的官方 PyTorch 实现：

用于图像重建与合成的焦点频率损失
Liming Jiang、Bo Dai、Wayne Wu 和 Chen Change Loy
发表于 ICCV 2021。
项目页面 | 论文 | 海报 | 幻灯片 | YouTube 演示

摘要： 得益于生成模型的发展，图像重建与合成取得了显著进展。然而，真实图像与生成图像之间仍可能存在差距，尤其是在频域方面。本研究证明，在频域中缩小这些差距可以进一步提升图像重建和合成的质量。我们提出了一种新颖的焦点频率损失函数，它能够通过降低易合成频率成分的权重，使模型自适应地关注那些难以合成的频率成分。该目标函数与现有的空间域损失互补，有效抑制了由于神经网络固有偏差而导致的重要频率信息丢失。我们展示了焦点频率损失在提升 VAE、pix2pix 和 SPADE 等流行模型的感知质量和定量性能方面的通用性和有效性，并进一步验证了其在 StyleGAN2 上的应用潜力。

更新

• [2021年9月] 焦点频率损失的 代码 已 发布。

• [2021年7月] 焦点频率损失的 论文 被 ICCV 2021 接收。

快速入门

运行 pip install focal-frequency-loss 进行安装。之后，您只需以下代码即可使用。

from focal_frequency_loss import FocalFrequencyLoss as FFL
ffl = FFL(loss_weight=1.0, alpha=1.0)  # 初始化 nn.Module 类

import torch
fake = torch.randn(4, 3, 64, 64)  # 替换为形状为 (N, C, H, W) 的预测张量
real = torch.randn(4, 3, 64, 64)  # 替换为形状为 (N, C, H, W) 的目标张量

loss = ffl(fake, real)  # 计算焦点频率损失

提示：

1. 当前支持的 PyTorch 版本：torch>=1.1.0。警告可忽略。请注意，论文中的实验是在 torch<=1.7.1,>=1.1.0 下进行的。
2. 初始化 FocalFrequencyLoss 类的参数：
   • loss_weight (float)：焦点频率损失的权重。默认值：1.0
   • alpha (float)：频谱权重矩阵的缩放因子，用于灵活性调整。默认值：1.0
   • patch_factor (int)：用于基于补丁的焦点频率损失时裁剪图像补丁的因子。默认值：1
   • ave_spectrum (bool)：是否使用小批量平均频谱。默认值：False
   • log_matrix (bool)：是否通过对数调整频谱权重矩阵。默认值：False
   • batch_matrix (bool)：是否使用基于批次的统计信息计算频谱权重矩阵。默认值：False
3. 经验表明，您需要调整的主要超参数是 loss_weight 和 alpha。通常首先需要调整损失权重。随后，较大的 alpha 值表示模型更加专注于高频成分。我们默认使用 alpha=1.0。

示例：图像重建（Vanilla AE）

作为指南，我们提供了一个在 CelebA 数据集上应用所提出的焦点频率损失（FFL）进行 Vanilla AE 图像重建的示例。应用 FFL 非常简单。核心细节请参见 此处。

安装

安装 Anaconda 后，建议您创建一个 Python 3.8.3 的新 conda 环境：

conda create -n ffl python=3.8.3 -y
conda activate ffl

克隆本仓库，安装 PyTorch 1.4.0（torch>=1.1.0 也可能适用）及其他依赖项：

git clone https://github.com/EndlessSora/focal-frequency-loss.git
cd focal-frequency-loss
pip install -r VanillaAE/requirements.txt

数据集准备

在此示例中，请从 CelebA 数据集的 官方网站 下载 img_align_celeba.zip。然后，强烈建议您解压该文件，并将 img_align_celeba 文件夹符号链接到 ./datasets/celeba，方法如下：

bash scripts/datasets/prepare_celeba.sh [PATH_TO_IMG_ALIGN_CELEBA]

或者您可以直接将 img_align_celeba 文件夹移动到 ./datasets/celeba。最终的目录结构应为：

├── datasets
│    ├── celeba
│    │    ├── img_align_celeba  
│    │    │    ├── 000001.jpg
│    │    │    ├── 000002.jpg
│    │    │    ├── 000003.jpg
│    │    │    ├── ...

测试与评估指标

下载 预训练模型 并将其解压到 ./VanillaAE/experiments。

我们提供了示例 测试脚本。如果您只有 CPU 环境，请在脚本中指定 --no_cuda。运行：

bash scripts/VanillaAE/test/celeba_recon_wo_ffl.sh
bash scripts/VanillaAE/test/celeba_recon_w_ffl.sh

Vanilla AE 图像重建的结果将默认保存在 ./VanillaAE/results 中。

测试完成后，您可以进一步计算本示例的评估指标。我们实现了一系列使用的 评估指标，并提供了 指标脚本。运行：

bash scripts/VanillaAE/metrics/celeba_recon_wo_ffl.sh
bash scripts/VanillaAE/metrics/celeba_recon_w_ffl.sh

您将看到各项指标的得分。指标日志将保存在 ./VanillaAE/results 中对应的实验文件夹内。

训练

我们提供了示例 训练脚本。如果您只有 CPU 环境，请在脚本中指定 --no_cuda。运行：

bash scripts/VanillaAE/train/celeba_recon_wo_ffl.sh
bash scripts/VanillaAE/train/celeba_recon_w_ffl.sh 

训练完成后，对新训练模型的推理过程与 测试与评估指标 类似。在 NVIDIA Tesla V100 GPU 上，使用 torch<=1.7.1,>=1.1.0 可以更好地复现结果。

更多结果

在这里，我们展示了在不同设置下应用所提出的焦点频率损失（FFL）的其他示例。

图像重建（VAE）

图像到图像转换（pix2pix | SPADE）

无条件图像合成（StyleGAN2）

256×256 分辨率的结果（未使用截断技巧）以及调整对比度后的迷你批次平均频谱图：

1024×1024 分辨率的结果（未使用截断技巧），由 StyleGAN2 结合 FFL 合成：

引用

如果您认为本工作对您的研究有所帮助，请引用我们的论文：

@inproceedings{jiang2021focal,
  title={Focal Frequency Loss for Image Reconstruction and Synthesis},
  author={Jiang, Liming and Dai, Bo and Wu, Wayne and Loy, Chen Change},
  booktitle={ICCV},
  year={2021}
}

致谢

Vanilla AE 的代码灵感来源于 PyTorch DCGAN 和 MUNIT。部分评估指标代码借用了 MMEditing。此外，我们还使用了 LPIPS 和 pytorch-fid 作为评估指标。

许可证

版权所有。代码以 MIT 许可证 发布。

版权所有 © 2021

其他实现

[非官方 TensorFlow 实现] 由 ZohebAbai 提供

Focal Frequency Loss 快速上手指南

Focal Frequency Loss (FFL) 是一种用于图像重建和合成的损失函数，通过自适应地关注难以合成的频率分量，弥补生成图像与真实图像在频域上的差距。该工具基于 PyTorch 实现，可轻松集成到 VAE、pix2pix、StyleGAN2 等模型中。

环境准备

• 操作系统：Linux / macOS / Windows
• Python 版本：推荐 Python 3.8+
• PyTorch 版本：torch >= 1.1.0（论文实验基于 1.1.0 <= torch <= 1.7.1）
• 依赖项：仅需安装 focal-frequency-loss 包，其他依赖由包自动管理

💡 国内用户建议使用清华或阿里镜像源加速安装。

安装步骤

使用 pip 直接安装官方发布包：

pip install focal-frequency-loss -i https://pypi.tuna.tsinghua.edu.cn/simple

若需从源码安装（可选）：

git clone https://github.com/EndlessSora/focal-frequency-loss.git
cd focal-frequency-loss
pip install -e .

基本使用

以下是最小可用示例，展示如何在训练循环中调用 FFL：

from focal_frequency_loss import FocalFrequencyLoss as FFL
ffl = FFL(loss_weight=1.0, alpha=1.0)  # 初始化损失模块

import torch
fake = torch.randn(4, 3, 64, 64)  # 模型预测输出 (N, C, H, W)
real = torch.randn(4, 3, 64, 64)  # 真实目标图像 (N, C, H, W)

loss = ffl(fake, real)  # 计算频域焦点损失
loss.backward()  # 反向传播

核心参数说明

• loss_weight (float): 损失权重，默认 1.0，通常需根据任务调整。
• alpha (float): 频谱权重矩阵的缩放因子，控制对难合成频率的关注程度，默认 1.0。值越大，模型越聚焦于高频细节。
• 其他高级参数（如 patch_factor, ave_spectrum 等）可按需配置，一般场景使用默认值即可。

✅ 建议优先调整 loss_weight，再微调 alpha 以获得最佳效果。