# easy-unet: Modular UNet Backbone in PyTorch
A **lightweight and flexible PyTorch library** providing a **modular UNet implementation** with advanced attention and normalization blocks. Designed for fast experimentation and development across diverse image processing and deep learning tasks.
## 🚀 Features
- 🧱 **UNet backbone** with residual blocks and flexible channel multipliers
- 🎯 **Advanced attention modules** including linear and flash attention for improved feature modeling
- ⚙️ **Configurable architecture** for dropout, channels, and dimensions
- 🧪 Modular and clean PyTorch codebase suitable for research and production
- 🔄 Supports easy integration with diffusion models, segmentation, or any custom pipeline
## 📦 Installation
```bash
pip install easy-unet
```
## 📁 Project Structure
```bash
easy-unet/
├── easy_unet/
│ ├── __init__.py
│ ├── module.py # All architecture classes and logic
├── pyproject.toml
├── LICENSE
└── README.md
```
## 🚀 Quick Start
### Import and create the model
```python
import torch
from easy_unet import UNet
model = UNet(
dim=64,
dim_mults=(1, 2, 4, 8),
channels=3,
dropout=0.1
)
x = torch.randn(1, 3, 256, 256) # sample input
output = model(x)
print(output.shape) # e.g., torch.Size([1, 3, 256, 256])
```
## ⚙️ Configuration Options
| Argument | Type | Default | Description |
|--|--|--|--|
| `dim` | `int` | `64` | Base number of feature channels |
| `dim_mults` | `tuple` | `(1, 2, 4, 8)` | Channel multipliers per U-Net stage |
| `channels` | `int` | `3` | Number of input/output channels (e.g., 3 for RGB) |
| `dropout` | `float` | `0.0` | Dropout rate for regularization. |
## 🙋♂️ Author
Developed by [Mehran Bazrafkan](mailto:mhrn.bzrafkn.dev@gmail.com)
> Created for general-purpose use cases and research requiring flexible UNet architectures in PyTorch.
## ⭐️ Support & Contribute
If you find this project useful, please:
- ⭐️ Star the repo
- 🐛 Report issues
- 📦 Suggest features or improvements
## 🔗 Related Projects
- [diffusion-pytorch-lib · PyPI (by me)](https://pypi.org/project/diffusion-pytorch-lib/)
- [variational-autoencoder-pytorch-lib · PyPI (by me)](https://pypi.org/project/variational-autoencoder-pytorch-lib/)
- [convolutional-autoencoder-pytorch · PyPI (by me)](https://pypi.org/project/convolutional-autoencoder-pytorch/)
## 📜 License
This project is licensed under the terms of the [`MIT LICENSE`](https://github.com/MehranBazrafkan/easy-unet/blob/main/LICENSE).
Raw data
{
"_id": null,
"home_page": null,
"name": "easy-unet",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.10",
"maintainer_email": null,
"keywords": "UNet, PyTorch, deep learning, image segmentation, attention, normalization, machine learning, computer vision",
"author": null,
"author_email": "Mehran Bazrafkan <mhrn.bzrafkn.dev@gmail.com>",
"download_url": "https://files.pythonhosted.org/packages/0c/24/abb91d56b6e0220b901329f753520c92795f7a91e23eff0996cd812d8191/easy_unet-0.1.0.tar.gz",
"platform": null,
"description": "\r\n# easy-unet: Modular UNet Backbone in PyTorch\r\n\r\nA **lightweight and flexible PyTorch library** providing a **modular UNet implementation** with advanced attention and normalization blocks. Designed for fast experimentation and development across diverse image processing and deep learning tasks.\r\n\r\n## \ud83d\ude80 Features\r\n\r\n- \ud83e\uddf1 **UNet backbone** with residual blocks and flexible channel multipliers \r\n- \ud83c\udfaf **Advanced attention modules** including linear and flash attention for improved feature modeling \r\n- \u2699\ufe0f **Configurable architecture** for dropout, channels, and dimensions \r\n- \ud83e\uddea Modular and clean PyTorch codebase suitable for research and production \r\n- \ud83d\udd04 Supports easy integration with diffusion models, segmentation, or any custom pipeline \r\n\r\n## \ud83d\udce6 Installation\r\n\r\n```bash\r\npip install easy-unet\r\n\r\n```\r\n\r\n## \ud83d\udcc1 Project Structure\r\n\r\n```bash\r\neasy-unet/\r\n\u251c\u2500\u2500 easy_unet/\r\n\u2502 \u251c\u2500\u2500 __init__.py\r\n\u2502 \u251c\u2500\u2500 module.py # All architecture classes and logic\r\n\u251c\u2500\u2500 pyproject.toml\r\n\u251c\u2500\u2500 LICENSE\r\n\u2514\u2500\u2500 README.md\r\n\r\n```\r\n\r\n## \ud83d\ude80 Quick Start\r\n\r\n### Import and create the model\r\n\r\n```python\r\nimport torch\r\nfrom easy_unet import UNet\r\n\r\nmodel = UNet(\r\n dim=64,\r\n dim_mults=(1, 2, 4, 8),\r\n channels=3,\r\n dropout=0.1\r\n)\r\n\r\nx = torch.randn(1, 3, 256, 256) # sample input\r\noutput = model(x)\r\nprint(output.shape) # e.g., torch.Size([1, 3, 256, 256])\r\n\r\n```\r\n\r\n## \u2699\ufe0f Configuration Options\r\n\r\n| Argument | Type | Default | Description |\r\n|--|--|--|--|\r\n| `dim` | `int` | `64` | Base number of feature channels |\r\n| `dim_mults` | `tuple` | `(1, 2, 4, 8)` | Channel multipliers per U-Net stage |\r\n| `channels` | `int` | `3` | Number of input/output channels (e.g., 3 for RGB) |\r\n| `dropout` | `float` | `0.0` | Dropout rate for regularization. |\r\n\r\n## \ud83d\ude4b\u200d\u2642\ufe0f Author\r\n\r\nDeveloped by [Mehran Bazrafkan](mailto:mhrn.bzrafkn.dev@gmail.com)\r\n\r\n> Created for general-purpose use cases and research requiring flexible UNet architectures in PyTorch.\r\n\r\n## \u2b50\ufe0f Support & Contribute\r\n\r\nIf you find this project useful, please:\r\n\r\n- \u2b50\ufe0f Star the repo\r\n\r\n- \ud83d\udc1b Report issues\r\n\r\n- \ud83d\udce6 Suggest features or improvements\r\n\r\n## \ud83d\udd17 Related Projects\r\n\r\n- [diffusion-pytorch-lib \u00b7 PyPI (by me)](https://pypi.org/project/diffusion-pytorch-lib/)\r\n- [variational-autoencoder-pytorch-lib \u00b7 PyPI (by me)](https://pypi.org/project/variational-autoencoder-pytorch-lib/)\r\n- [convolutional-autoencoder-pytorch \u00b7 PyPI (by me)](https://pypi.org/project/convolutional-autoencoder-pytorch/)\r\n\r\n## \ud83d\udcdc License\r\n\r\nThis project is licensed under the terms of the [`MIT LICENSE`](https://github.com/MehranBazrafkan/easy-unet/blob/main/LICENSE).\r\n",
"bugtrack_url": null,
"license": null,
"summary": "A lightweight and flexible PyTorch library providing a modular UNet implementation with advanced attention and normalization blocks for fast image processing and deep learning development.",
"version": "0.1.0",
"project_urls": null,
"split_keywords": [
"unet",
" pytorch",
" deep learning",
" image segmentation",
" attention",
" normalization",
" machine learning",
" computer vision"
],
"urls": [
{
"comment_text": null,
"digests": {
"blake2b_256": "25e2d566bc133a7d1e02329c7033408f7092545e58640a7226078e1524d828f4",
"md5": "7aa7537f1d0f6ed6fff1956d4295ca6f",
"sha256": "caec071e599e12337309bd297c1c8bf0e2482775881f994e01974c0b8dfbe89d"
},
"downloads": -1,
"filename": "easy_unet-0.1.0-py3-none-any.whl",
"has_sig": false,
"md5_digest": "7aa7537f1d0f6ed6fff1956d4295ca6f",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.10",
"size": 6311,
"upload_time": "2025-08-11T09:47:36",
"upload_time_iso_8601": "2025-08-11T09:47:36.587596Z",
"url": "https://files.pythonhosted.org/packages/25/e2/d566bc133a7d1e02329c7033408f7092545e58640a7226078e1524d828f4/easy_unet-0.1.0-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": null,
"digests": {
"blake2b_256": "0c24abb91d56b6e0220b901329f753520c92795f7a91e23eff0996cd812d8191",
"md5": "2fdfc7ef05a8fcf8faeb5e0fa4735a17",
"sha256": "6eb3196e1b7ce506542f02cc5202513a0d030bb7998255308091ba75d8c85f7c"
},
"downloads": -1,
"filename": "easy_unet-0.1.0.tar.gz",
"has_sig": false,
"md5_digest": "2fdfc7ef05a8fcf8faeb5e0fa4735a17",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.10",
"size": 5964,
"upload_time": "2025-08-11T09:47:38",
"upload_time_iso_8601": "2025-08-11T09:47:38.049632Z",
"url": "https://files.pythonhosted.org/packages/0c/24/abb91d56b6e0220b901329f753520c92795f7a91e23eff0996cd812d8191/easy_unet-0.1.0.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2025-08-11 09:47:38",
"github": false,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"lcname": "easy-unet"
}
JSON