Uformer - Pytorch
Implementation of Uformer, Attention-based Unet, in Pytorch. It will only offer the concat-cross-skip connection.
This repository will be geared towards use in a project for learning protein structures. Specifically, it will include the ability to condition on time steps (needed for DDPM), as well as 2d relative positional encoding using rotary embeddings (instead of the bias on the attention matrix in the paper).
Install
$ pip install uformer-pytorch
Usage
import torch
from uformer_pytorch import Uformer
model = Uformer(
dim = 64, # initial dimensions after input projection, which increases by 2x each stage
stages = 4, # number of stages
num_blocks = 2, # number of transformer blocks per stage
window_size = 16, # set window size (along one side) for which to do the attention within
dim_head = 64,
heads = 8,
ff_mult = 4
)
x = torch.randn(1, 3, 256, 256)
pred = model(x) # (1, 3, 256, 256)
To condition on time for DDPM training
import torch
from uformer_pytorch import Uformer
model = Uformer(
dim = 64,
stages = 4,
num_blocks = 2,
window_size = 16,
dim_head = 64,
heads = 8,
ff_mult = 4,
time_emb = True # set this to true
)
x = torch.randn(1, 3, 256, 256)
time = torch.arange(1)
pred = model(x, time = time) # (1, 3, 256, 256)
Citations
@misc{wang2021uformer,
title = {Uformer: A General U-Shaped Transformer for Image Restoration},
author = {Zhendong Wang and Xiaodong Cun and Jianmin Bao and Jianzhuang Liu},
year = {2021},
eprint = {2106.03106},
archivePrefix = {arXiv},
primaryClass = {cs.CV}
}
84 Nov 22, 2022
1 Oct 21, 2021
4 Dec 21, 2021
1.8k Jan 7, 2023
3 Jun 29, 2022
556 Jan 4, 2023
71 Dec 30, 2022
5 Nov 21, 2022
0 Sep 2, 2021
4 Jan 07, 2023
84 Jan 04, 2023
11 Oct 20, 2022
15 Oct 08, 2022
72 Oct 20, 2022
1.3k Dec 23, 2022
32 Nov 25, 2022
55 Dec 24, 2022
55 Dec 09, 2022
16 Dec 15, 2022
32 Nov 20, 2022
4 Aug 22, 2022
2.7k Dec 31, 2022
181 Dec 29, 2022
10 Jun 17, 2022
0 Oct 09, 2022
49 Dec 13, 2022
0 Jan 20, 2022
3 Feb 04, 2022
28 Dec 15, 2022