Implementation of Uformer, Attention-based Unet, in Pytorch

Last update: Dec 19, 2022

Overview

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}
}

You might also like...

Implementation detail for paper "Multi-level colonoscopy malignant tissue detection with adversarial CAC-UNet"

Multi-level-colonoscopy-malignant-tissue-detection-with-adversarial-CAC-UNet Implementation detail for our paper "Multi-level colonoscopy malignant ti

84 Nov 22, 2022

Implementation of UNet on the Joey ML framework

Independent Research Project - Code Joey can be cloned from here https://github.com/devitocodes/joey/. Devito and other dependencies such as PyTorch a

1 Oct 21, 2021

Implementation of UNET architecture for Image Segmentation.

Semantic Segmentation using UNET This is the implementation of UNET on Carvana Image Masking Kaggle Challenge About the Dataset This dataset contains

4 Dec 21, 2021

Official Keras Implementation for UNet++ in IEEE Transactions on Medical Imaging and DLMIA 2018

UNet++: A Nested U-Net Architecture for Medical Image Segmentation UNet++ is a new general purpose image segmentation architecture for more accurate i

1.8k Jan 7, 2023

A unet implementation for Image semantic segmentation

Unet-pytorch a unet implementation for Image semantic segmentation 参考网上的Unet做分割的代码，做了一个针对kaggle地盐识别的，请去以下地址获取数据集: https://www.kaggle.com/c/tgs-salt-id

3 Jun 29, 2022

RETRO-pytorch - Implementation of RETRO, Deepmind's Retrieval based Attention net, in Pytorch

RETRO - Pytorch (wip) Implementation of RETRO, Deepmind's Retrieval based Attent

556 Jan 4, 2023

PyTorch code for our paper "Attention in Attention Network for Image Super-Resolution"

Under construction... Attention in Attention Network for Image Super-Resolution (A2N) This repository is an PyTorch implementation of the paper "Atten

71 Dec 30, 2022

Unet network with mean teacher for altrasound image segmentation

5 Nov 21, 2022

Hippocampal segmentation using the UNet network for each axis

Hipposeg Hippocampal segmentation using the UNet network for each axis, inspired by https://github.com/MICLab-Unicamp/e2dhipseg Red: False Positive Gr

0 Sep 2, 2021

Implementation of Uformer, Attention-based Unet, in Pytorch

Related tags

Overview

Uformer - Pytorch

Install

Usage

Citations

You might also like...

Implementation detail for paper "Multi-level colonoscopy malignant tissue detection with adversarial CAC-UNet"

Implementation of UNet on the Joey ML framework

Implementation of UNET architecture for Image Segmentation.

Official Keras Implementation for UNet++ in IEEE Transactions on Medical Imaging and DLMIA 2018

A unet implementation for Image semantic segmentation

RETRO-pytorch - Implementation of RETRO, Deepmind's Retrieval based Attention net, in Pytorch

PyTorch code for our paper "Attention in Attention Network for Image Super-Resolution"

Unet network with mean teacher for altrasound image segmentation

Hippocampal segmentation using the UNet network for each axis

Releases(0.0.8)

0.0.8(Oct 26, 2021)

0.0.7(Aug 24, 2021)

0.0.6(Jun 17, 2021)

0.0.5(Jun 17, 2021)

0.0.4(Jun 17, 2021)

0.0.3(Jun 17, 2021)

0.0.2(Jun 17, 2021)

0.0.1(Jun 17, 2021)

Owner

Phil Wang

PyTorch code for ICPR 2020 paper Future Urban Scene Generation Through Vehicle Synthesis

Building a real-time environment using webcam frame division in OpenCV and classify cropped images using a fine-tuned vision transformers on hybryd datasets samples for facial emotion recognition.

A fast MoE impl for PyTorch

Implementations of LSTM: A Search Space Odyssey variants and their training results on the PTB dataset.

An implementation of a sequence to sequence neural network using an encoder-decoder

A Simplied Framework of GAN Inversion

A pyparsing-based library for parsing SOQL statements

A really easy-to-use and powerful sudoku solver.

Code for training and evaluation of the model from "Language Generation with Recurrent Generative Adversarial Networks without Pre-training"

Pretrained models for Jax/Flax: StyleGAN2, GPT2, VGG, ResNet.

Large dataset storage format for Pytorch

Simple API for UCI Machine Learning Dataset Repository (search, download, analyze)

Auto White-Balance Correction for Mixed-Illuminant Scenes

OCTIS: Comparing Topic Models is Simple! A python package to optimize and evaluate topic models (accepted at EACL2021 demo track)

This is a deep learning-based method to segment deep brain structures and a brain mask from T1 weighted MRI.

Semantic Segmentation with SegFormer on Drone Dataset.

Video Background Music Generation with Controllable Music Transformer (ACM MM 2021 Oral)

Deep deconfounded recommender (Deep-Deconf) for paper "Deep causal reasoning for recommendations"

Robustness between the worst and average case

Deep Dual Consecutive Network for Human Pose Estimation (CVPR2021)