注意力网络中的注意力图像超分辨率模型（A2N）

A2N模型在DIV2K（800张训练图像，扩充到4000张图像，100张验证图像）上进行了预训练，用于2倍、3倍和4倍的图像超分辨率。它是由Chen等人在2021年的论文中引入的，首次在 this repository 中发布。

图像超分辨率的目标是从单个低分辨率（LR）图像恢复高分辨率（HR）图像。下图显示了真实图像（HR）、双三次上采样x2和模型上采样x2。

模型描述

A2N模型提出了一种注意力嵌套网络（A2N），用于高精度图像超分辨率。具体而言，A2N由非注意力分支和耦合注意力分支组成。提出了注意力丢弃模块，用于根据输入特征生成动态注意力权重，以抑制不需要的注意力调整。这使得注意力模块能够专门处理有益的示例，而不会受到其他惩罚，从而大大提高了注意力网络的容量，而参数开销很小。

更重要的是，该模型具有轻量级和快速训练的特点（约1.5m个参数，约4mb）。

预期用途和限制

您可以使用预训练模型将图像放大2倍、3倍和4倍。您还可以使用训练器对自己的数据集进行模型训练。

如何使用

该模型可以与 super_image 库一起使用：

pip install super-image

以下是如何使用预训练模型将图像放大的示例：

from super_image import A2nModel, ImageLoader
from PIL import Image
import requests

url = 'https://paperswithcode.com/media/datasets/Set5-0000002728-07a9793f_zA3bDjj.jpg'
image = Image.open(requests.get(url, stream=True).raw)

model = A2nModel.from_pretrained('eugenesiow/a2n', scale=2)      # scale 2, 3 and 4 models available
inputs = ImageLoader.load_image(image)
preds = model(inputs)

ImageLoader.save_image(preds, './scaled_2x.png')                        # save the output 2x scaled image to `./scaled_2x.png`
ImageLoader.save_compare(inputs, preds, './scaled_2x_compare.png')      # save an output comparing the super-image with a bicubic scaling

训练数据

2倍、3倍和4倍图像超分辨率模型是在 DIV2K 上进行预训练的，该数据集包含800张高质量（2K分辨率）的训练图像，扩充到4000张图像，并使用100张验证图像（图像编号为801到900）。

训练过程

预处理

我们遵循 Wang et al. 的预处理和训练方法。使用双三次插值作为调整大小的方法，将高分辨率（HR）图像的大小降低2倍、3倍和4倍，创建低分辨率（LR）图像。在训练过程中，使用从LR输入中提取的大小为64×64的RGB块以及其相应的HR块。在预处理阶段对训练集应用数据增强，在原始图像的四个角和中心创建五张图像。

我们需要huggingface datasets 库来下载数据：

pip install datasets

以下代码获取数据并预处理/增强数据。

from datasets import load_dataset
from super_image.data import EvalDataset, TrainDataset, augment_five_crop

augmented_dataset = load_dataset('eugenesiow/Div2k', 'bicubic_x4', split='train')\
    .map(augment_five_crop, batched=True, desc="Augmenting Dataset")                                # download and augment the data with the five_crop method
train_dataset = TrainDataset(augmented_dataset)                                                     # prepare the train dataset for loading PyTorch DataLoader
eval_dataset = EvalDataset(load_dataset('eugenesiow/Div2k', 'bicubic_x4', split='validation'))      # prepare the eval dataset for the PyTorch DataLoader

预训练

模型在GPU上进行了训练。下面提供了训练代码：

from super_image import Trainer, TrainingArguments, A2nModel, A2nConfig

training_args = TrainingArguments(
    output_dir='./results',                 # output directory
    num_train_epochs=1000,                  # total number of training epochs
)

config = A2nConfig(
    scale=4,                                # train a model to upscale 4x
)
model = A2nModel(config)

trainer = Trainer(
    model=model,                         # the instantiated model to be trained
    args=training_args,                  # training arguments, defined above
    train_dataset=train_dataset,         # training dataset
    eval_dataset=eval_dataset            # evaluation dataset
)

trainer.train()

评估结果

评估指标包括 PSNR 和 SSIM 。

评估数据集包括：

• Set5 - Bevilacqua et al. (2012)
• Set14 - Zeyde et al. (2010)
• BSD100 - Martin et al. (2001)
• Urban100 - Huang et al. (2015)

以下结果列用PSNR/SSIM表示，与双三次插值作为基准进行比较。

|数据集|缩放|双三次插值|A2N||---|---|---|---||Set5|2x|33.64/0.9292|37.87/0.9602||Set5|3x|30.39/0.8678|34.8/0.9387||Set5|4x|28.42/0.8101|32.07/0.8933||Set14|2x|30.22/0.8683|33.45/0.9162||Set14|3x|27.53/0.7737|30.94/0.8568||Set14|4x|25.99/0.7023|28.56/0.7801||BSD100|2x|29.55/0.8425|32.11/0.8987||BSD100|3x|27.20/0.7382|29.56/0.8173||BSD100|4x|25.96/0.6672|27.54/0.7342||Urban100|2x|26.66/0.8408|31.71/0.9240||Urban100|3x||28.95/0.8671||Urban100|4x|23.14/0.6573|25.89/0.7787|

您可以在以下链接找到一个笔记本，可以轻松运行对预训练模型的评估：

BibTeX引用和引文信息

@misc{chen2021attention,
      title={Attention in Attention Network for Image Super-Resolution}, 
      author={Haoyu Chen and Jinjin Gu and Zhi Zhang},
      year={2021},
      eprint={2104.09497},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}