U-Shape Networks Simple Code

A lightweight and easy-to-understand PyTorch implementation of U-shaped encoder–decoder networks.

This repository is designed for rapid experimentation, architecture comparison, and educational use. It provides simplified implementations of three backbone variants based on NAFNet, ResNet with CBAM, and Restormer Transformer, while preserving the main components of a U-shaped network, including encoder–decoder stages, skip connections, downsampling, upsampling, and global residual learning.

Requirements

pip install torch thop

Run

python net.py

The profiling script reports:

Input shape
Output shape
Parameters
MACs
FLOPs

References

1. L. Chen, X. Chu, X. Zhang, and J. Sun, “Simple Baselines for Image Restoration,” European Conference on Computer Vision (ECCV), 2022.
   https://arxiv.org/abs/2204.04676

2. S. Woo, J. Park, J.-Y. Lee, and I. S. Kweon, “CBAM: Convolutional Block Attention Module,” European Conference on Computer Vision (ECCV), 2018.
   https://arxiv.org/abs/1807.06521

3. S. W. Zamir et al., “Restormer: Efficient Transformer for High-Resolution Image Restoration,” IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
   https://arxiv.org/abs/2111.09881

4. Lyken17, “THOP: PyTorch-OpCounter.”
   https://github.com/Lyken17/pytorch-OpCounter

MIT

This is free, get me star if u like