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

Multi-Scale Geometric Consistency Guided Multi-View Stereo

[cvpr22] Perturbed and Strict Mean Teachers for Semi-supervised Semantic Segmentation

Source code and dataset for ACL2021 paper: "ERICA: Improving Entity and Relation Understanding for Pre-trained Language Models via Contrastive Learning".

Code repository for our paper regarding the L3D dataset.

Flexible Networks for Learning Physical Dynamics of Deformable Objects (2021)

Attention-based CNN-LSTM and XGBoost hybrid model for stock prediction

Code for C2-Matching (CVPR2021). Paper: Robust Reference-based Super-Resolution via C2-Matching.

K-Means Clustering and Hierarchical Clustering Unsupervised Learning Solution in Python3.

Spiking Neural Network for Computer Vision using SpikingJelly framework and Pytorch-Lightning

I created My own Virtual Artificial Intelligence named genesis, He can assist with my Tasks and also perform some analysis,,

The code for two papers: Feedback Transformer and Expire-Span.

Moving Object Segmentation in 3D LiDAR Data: A Learning-based Approach Exploiting Sequential Data

Code for IntraQ, PyTorch implementation of our paper under review

PyTorch reimplementation of the paper Involution: Inverting the Inherence of Convolution for Visual Recognition [CVPR 2021].

Library of deep learning models and datasets designed to make deep learning more accessible and accelerate ML research.

Pytorch implementation for "Distribution-Balanced Loss for Multi-Label Classification in Long-Tailed Datasets" (ECCV 2020 Spotlight)

Lunar is a neural network aimbot that uses real-time object detection accelerated with CUDA on Nvidia GPUs.

The code for SAG-DTA: Prediction of Drug–Target Affinity Using Self-Attention Graph Network.

InDuDoNet+: A Model-Driven Interpretable Dual Domain Network for Metal Artifact Reduction in CT Images

Hcaptcha-challenger - Gracefully face hCaptcha challenge with Yolov5(ONNX) embedded solution