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

Generative vs Discriminative: Rethinking The Meta-Continual Learning (NeurIPS 2021)

HiddenMarkovModel implements hidden Markov models with Gaussian mixtures as distributions on top of TensorFlow

DSL for matching Python ASTs

Opinionated code formatter, just like Python's black code formatter but for Beancount

Covid19-Forecasting - An interactive website that tracks, models and predicts COVID-19 Cases

List of papers, code and experiments using deep learning for time series forecasting

Code and data for ACL2021 paper Cross-Lingual Abstractive Summarization with Limited Parallel Resources.

Pytorch Implementation of "Desigining Network Design Spaces", Radosavovic et al. CVPR 2020.

Image morphing without reference points by applying warp maps and optimizing over them.

A python/pytorch utility library

In the case of your data having only 1 channel while want to use timm models

Official Implementation for HyperStyle: StyleGAN Inversion with HyperNetworks for Real Image Editing

Code for KHGT model, AAAI2021

TensorFlow 101: Introduction to Deep Learning for Python Within TensorFlow

[CVPR 2022 Oral] Versatile Multi-Modal Pre-Training for Human-Centric Perception

Densely Connected Search Space for More Flexible Neural Architecture Search (CVPR2020)

Implementation of the "Point 4D Transformer Networks for Spatio-Temporal Modeling in Point Cloud Videos" paper.

LoveDA: A Remote Sensing Land-Cover Dataset for Domain Adaptive Semantic Segmentation

PyTorch-lightning implementation of the ESFW module proposed in our paper Edge-Selective Feature Weaving for Point Cloud Matching

codes for IKM (arXiv2021, Submitted to IEEE Trans)