HAT

在图像超分辨率变换器中激活更多像素 [论文链接]

陈向宇、王新涛、周建涛、乔宇和董超

HAT：用于图像复原的混合注意力变换器 [论文链接]

陈向宇、王新涛、张文龙、孔祥涛、周建涛、乔宇和董超

更新

✅ 2022-05-09：在Arxiv上发布论文的第一个版本。
✅ 2022-05-20：发布HAT的代码、模型和结果。
✅ 2022-08-29：为SRx4添加Replicate演示。
✅ 2022-09-25：为有限GPU内存的推理添加分块模式。
✅ 2022-11-24：上传用于现实世界超分辨率的基于GAN的HAT模型（Real_HAT_GAN_SRx4.pth）。
✅ 2023-03-19：更新论文至CVPR版本。添加小型HAT模型。
✅ 2023-04-05：上传HAT-S的代码、模型和结果。
✅ 2023-08-01：上传另一个GAN模型以获得更清晰的结果（Real_HAT_GAN_SRx4_sharper.pth）。
✅ 2023-08-01：上传现实世界基于GAN模型的训练配置。
✅ 2023-09-11：在Arxiv上发布论文的扩展版本。
**（待办）**为Replicate演示添加分块模式。
**（待办）**更新Replicate演示以支持现实世界超分辨率。
**（待办）**为多种图像复原任务添加HAT模型。

概述

**在不使用ImageNet预训练的情况下，SRx4的基准测试结果。Multi-Adds是针对64x64输入计算的。**

模型	参数量(M)	Multi-Adds(G)	Set5	Set14	BSD100	Urban100	Manga109
SwinIR	11.9	53.6	32.92	29.09	27.92	27.45	32.03
HAT-S	9.6	54.9	32.92	29.15	27.97	27.87	32.35
HAT	20.8	102.4	33.04	29.23	28.00	27.97	32.48

真实世界超分辨率结果

注意：

训练配置中的默认设置（几乎与Real-ESRGAN相同）用于训练Real_HAT_GAN_SRx4_sharper。
Real_HAT_GAN_SRx4使用类似的设置进行训练，但不对真实图像进行USM处理。
Real_HAT_GAN_SRx4会有更好的保真度。
Real_HAT_GAN_SRx4_sharper会有更好的感知质量。

以下结果由Real_HAT_GAN_SRx4_sharper.pth生成。

与最先进的真实超分辨率方法的比较。

引用

BibTeX

@InProceedings{chen2023activating,
    author    = {Chen, Xiangyu and Wang, Xintao and Zhou, Jiantao and Qiao, Yu and Dong, Chao},
    title     = {Activating More Pixels in Image Super-Resolution Transformer},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2023},
    pages     = {22367-22377}
}

@article{chen2023hat,
  title={HAT: Hybrid Attention Transformer for Image Restoration},
  author={Chen, Xiangyu and Wang, Xintao and Zhang, Wenlong and Kong, Xiangtao and Qiao, Yu and Zhou, Jiantao and Dong, Chao},
  journal={arXiv preprint arXiv:2309.05239},
  year={2023}
}

环境

PyTorch >= 1.7 （强烈建议不要使用torch 1.8！！！它会导致异常性能。）
BasicSR == 1.3.4.9

安装

首先安装PyTorch。然后，

pip install -r requirements.txt
python setup.py develop

如何测试

无需实现代码，chaiNNer是运行我们模型的一个不错的工具。

否则，

参考./options/test中要测试的模型的配置文件，准备测试数据和预训练模型。
预训练模型可在Google Drive或百度网盘（提取码：qyrl）获取。
然后运行以下代码（以HAT_SRx4_ImageNet-pretrain.pth为例）：

python hat/test.py -opt options/test/HAT_SRx4_ImageNet-pretrain.yml

测试结果将保存在./results文件夹中。

参考./options/test/HAT_SRx4_ImageNet-LR.yml进行无真实图像的推理。

注意，在测试时也提供了分块模式，用于GPU内存有限的情况。您可以参考./options/test/HAT_tile_example.yml来修改自定义测试选项中的具体分块设置。

如何训练

参考./options/train中要训练的模型的配置文件。
训练数据的准备可以参考此页面。ImageNet数据集可以在官方网站下载。
训练命令如下：

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m torch.distributed.launch --nproc_per_node=8 --master_port=4321 hat/train.py -opt options/train/train_HAT_SRx2_from_scratch.yml --launcher pytorch

请注意，每个GPU的默认批量大小为4，这将为每个GPU消耗约20G内存。

训练日志和权重将保存在./experiments文件夹中。

结果

基准数据集的推理结果可在Google Drive或百度网盘（访问码：63p5）获取。

联系方式

如有任何疑问，请发送邮件至chxy95@gmail.com，或加入BasicSR的微信群与作者讨论。