Hybrid Attention Transformer (HAT)

Paper

HAT: Hybrid Attention Transformer for Image Restoration

Architecture Overview

The HAT model consists of three main components: shallow feature extraction, deep feature extraction, and image reconstruction.

Algorithm Principle

The HAT approach integrates channel attention and window-based self-attention mechanisms, leveraging their complementary strengths. Additionally, an overlapping cross-attention module is introduced to enhance interaction between neighboring windows, thereby improving cross-window information aggregation. During trainign, HAT employs a shared task pre-training strategy to further exploit the model's potential for performance enhancement. These design choices enable HAT to activate more pixels during reconstruction, resulting in significant performance improvements.

Environment Setup

-v Adjust paths, docker_name, and imageID according to your environment

Docker (Method One)

docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:1.13.1-centos7.6-dtk23.10-py38

docker run -it -v /path/your_code_data/:/path/your_code_data/ -v /opt/hyhal/:/opt/hyhal/ --shm-size=32G --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video --name docker_name imageID bash

cd /your_code_path/hat_pytorch
pip install -r requirements.txt
python setup.py develop

Dockerfile (Method Two)

cd ./docker
cp ../requirements.txt requirements.txt

docker build --no-cache -t hat:latest .
docker run -it -v /path/your_code_data/:/path/your_code_data/ -v /opt/hyhal/:/opt/hyhal/ --shm-size=32G --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video --name docker_name imageID bash

cd /your_code_path/hat_pytorch
python setup.py develop

Anaconda (Method Three)

1. Special deep learning libraries required for DCU GPUs can be downloaded from the HPCCube Developer Community: https://developer.hpccube.com/tool/

DTK Stack: dtk23.10
Python: python3.8
Torch: 1.13.1
TorchVision: 0.14.1

Note: The versions of DTK stack, Python, Torch, and other DCU-related tools must correspond exactly.

2. Install other non-specialized libraries via requirements.txt

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

Dataset

Training: ImageNet dataset DIV2K Flickr2K

Note: DF2K combines DIV2K and Flickr2K datasets. Refer to BasicSR for data processing.

Testing: Classical SR Testing

Dataset preparation steps:

1. Store data under the datasets directory.
2. For BSD100 and Urban100, create new directories GTmod4 and LRbicx4 inside each respective folder, then place original images into GTmod4. Execute the following commands from the datasets directory:

python gen_LRbicx4.py --file_name ./BSD100
python gen_LRbicx4.py --file_name ./urban100

1. Expected dataset directory structure:

├── DF2K
│   ├── DF2K_HR # High-resolution data
│   ├── DF2K_HR_sub # Generated patches
│   ├── DF2K_bicx4 # Training low-resolution bicubic X4 data
│   ├── DF2K_bicx4_sub # Generated patches
├── Set5
│   ├── GTmod12
│   ├── LRbicx2
│   ├── LRbicx3
│   ├── LRbicx4
├── Set14
│   ├── GTmod12
│   ├── LRbicx2
│   ├── LRbicx3
│   ├── LRbicx4
├── BSDS100
│   ├── GTmod4 # Original images
│   ├── LRbicx4
├── urban100
│   ├── GTmod4 # Original images
│   ├── LRbicx4

Note: A tiny_datasets is available for quick start; replace relevant paths if using this small dataset instead of full dataset.

1. Since the DF2K dataset is at 2K resolution (e.g., 2048x1080), it's often cropped into overlapping sub-images of 480x480 for training, which are then randomly cropped into patches like 128x128 or 192x192.

python extract_subimages.py # Crop images into sub-images

1. Generate metadata files

python generate_meta_info.py

Training

Pre-trained models can be downloaded from Google Drive or Baidu Pan (access code: qyrl).

Training logs and weights will be saved in the ./experiments directory.

Single Machine, Multiple GPUs

# Default config: train_HAT_SRx4_finetune_from_ImageNet_pretrain.yml
bash train.sh

Multi-Machine, Multiple GPUs

In multi-node scenarios, list nodes in hostfile, one per line (e.g., c1xxxxxx slots=4).

1. Modify virtual environment path in run_train_multi.sh at line 18;
2. Update YAML config file path in single_process.sh at line 22, unless default is used.

Execute the command below; logs will be stored in the logs folder.

# Default config: train_HAT_SRx4_finetune_from_ImageNet_pretrain.yml
bash run_train_multi.sh

Inference

Pre-trained models can be downloaded from Google Drive or Baidu Pan (access code: qyrl).

Results will be saved in the ./results directory. The option file options/test/HAT_SRx4_ImageNet-LR.yml is used for inference without ground truth images.

# Default config: HAT_SRx4_ImageNet-pretrain.yml
bash val.sh

Results

Test results based on Real_HAT_GAN_SRx4_sharper.pth

Accuracy Metrics

HAT

Model	Params(M)	Multi-Adds(G)	Set5	Set14	BSD100	Urban100
Z100L	20.8	102.4	33.1486	29.3587	25.4074	21.2687

Application Scenarios

Algorithm Category

Image Reconstruction

Industry Applications

Transportation, Public Security, Manufacturing

Source Repository & Feedback

• ModelZoo / HAT_pytorch · GitLab

References

• GitHub - XPixelGroup/HAT: CVPR2023 - Activating More Pixels in Image Super-Resolution Transformer Arxiv - HAT: Hybrid Attention Transformer for Image Restoration