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

RuDOLPH: One Hyper-Modal Transformer can be creative as DALL-E and smart as CLIP

Dynamic Attentive Graph Learning for Image Restoration, ICCV2021 [PyTorch Code]

This repository contains FEDOT - an open-source framework for automated modeling and machine learning (AutoML)

A DCGAN to generate anime faces using custom mined dataset

Speed-Test - You can check your intenet speed using this tool

Codebase for Time-series Generative Adversarial Networks (TimeGAN)

Hummingbird compiles trained ML models into tensor computation for faster inference.

Implementation of Shape Generation and Completion Through Point-Voxel Diffusion

Unofficial implementation of MLP-Mixer: An all-MLP Architecture for Vision

You can draw the corresponding bounding box into the image and save it according to the result file (txt format) run by the tracker.

Relative Uncertainty Learning for Facial Expression Recognition

MiniHack the Planet: A Sandbox for Open-Ended Reinforcement Learning Research

Constraint-based geometry sketcher for blender

A high-performance Python-based I/O system for large (and small) deep learning problems, with strong support for PyTorch.

the code for paper "Energy-Based Open-World Uncertainty Modeling for Confidence Calibration"

Get 2D point positions (e.g., facial landmarks) projected on 3D mesh

Deep Structured Instance Graph for Distilling Object Detectors (ICCV 2021)

Changing the Mind of Transformers for Topically-Controllable Language Generation

Learning to Reach Goals via Iterated Supervised Learning

This code is for eCaReNet: explainable Cancer Relapse Prediction Network.