https://github.com/lucidrains/ittr-pytorch
Implementation of the Hybrid Perception Block and Dual-Pruned Self-Attention block from the ITTR paper for Image to Image Translation using Transformers
https://github.com/lucidrains/ittr-pytorch
artificial-intelligence attention-mechanism deep-learning image-to-image-translation
Last synced: about 1 year ago
JSON representation
Implementation of the Hybrid Perception Block and Dual-Pruned Self-Attention block from the ITTR paper for Image to Image Translation using Transformers
- Host: GitHub
- URL: https://github.com/lucidrains/ittr-pytorch
- Owner: lucidrains
- License: mit
- Created: 2022-04-01T18:28:24.000Z (over 4 years ago)
- Default Branch: main
- Last Pushed: 2022-04-02T02:14:02.000Z (over 4 years ago)
- Last Synced: 2025-04-28T13:07:46.997Z (about 1 year ago)
- Topics: artificial-intelligence, attention-mechanism, deep-learning, image-to-image-translation
- Language: Python
- Homepage:
- Size: 146 KB
- Stars: 35
- Watchers: 3
- Forks: 8
- Open Issues: 3
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
## ITTR - Pytorch
Implementation of the Hybrid Perception Block (`HPB`) and Dual-Pruned Self-Attention (`DPSA`) block from the ITTR paper for Image to Image Translation using Transformers.
## Install
```bash
$ pip install ITTR-pytorch
```
## Usage
They had 9 blocks of Hybrid Perception Block (HPB) in the paper
```python
import torch
from ITTR_pytorch import HPB
block = HPB(
dim = 512, # dimension
dim_head = 32, # dimension per attention head
heads = 8, # number of attention heads
attn_height_top_k = 16, # number of top indices to select along height, for the attention pruning
attn_width_top_k = 16, # number of top indices to select along width, for the attention pruning
attn_dropout = 0., # attn dropout
ff_mult = 4, # expansion factor of feedforward
ff_dropout = 0. # feedforward dropout
)
fmap = torch.randn(1, 512, 32, 32)
out = block(fmap) # (1, 512, 32, 32)
```
You can also use the dual-pruned self-attention as so
```python
import torch
from ITTR_pytorch import DPSA
attn = DPSA(
dim = 512, # dimension
dim_head = 32, # dimension per attention head
heads = 8, # number of attention heads
height_top_k = 48, # number of top indices to select along height, for the attention pruning
width_top_k = 48, # number of top indices to select along width, for the attention pruning
dropout = 0. # attn dropout
)
fmap = torch.randn(1, 512, 32, 32)
out = attn(fmap) # (1, 512, 32, 32)
```
## Citations
```bibtex
@inproceedings{Zheng2022ITTRUI,
title = {ITTR: Unpaired Image-to-Image Translation with Transformers},
author = {Wanfeng Zheng and Qiang Li and Guoxin Zhang and Pengfei Wan and Zhongyuan Wang},
year = {2022}
}
```