# Zuko - Normalizing flows in PyTorch
Zuko is a Python package that implements normalizing flows in [PyTorch](https://pytorch.org). It relies as much as possible on distributions and transformations already provided by PyTorch. Unfortunately, the `Distribution` and `Transform` classes of `torch` are not sub-classes of `torch.nn.Module`, which means you cannot send their internal tensors to GPU with `.to('cuda')` or retrieve their parameters with `.parameters()`. Worse, the concepts of conditional distribution and transformation, which are essential for probabilistic inference, are impossible to express.
To solve these problems, `zuko` defines two concepts: the `LazyDistribution` and `LazyTransform`, which are any modules whose forward pass returns a `Distribution` or `Transform`, respectively. Because the creation of the actual distribution/transformation is delayed, an eventual condition can be easily taken into account. This design enables lazy distributions, including normalizing flows, to act like distributions while retaining features inherent to modules, such as trainable parameters. It also makes the implementations easy to understand and extend.
> In the [Avatar](https://wikipedia.org/wiki/Avatar:_The_Last_Airbender) cartoon, [Zuko](https://wikipedia.org/wiki/Zuko) is a powerful firebender 🔥
## Acknowledgements
Zuko takes significant inspiration from [nflows](https://github.com/bayesiains/nflows) and [Stefan Webb](https://github.com/stefanwebb)'s work in [Pyro](https://github.com/pyro-ppl/pyro) and [FlowTorch](https://github.com/facebookincubator/flowtorch).
## Installation
The `zuko` package is available on [PyPI](https://pypi.org/project/zuko), which means it is installable via `pip`.
```
pip install zuko
```
Alternatively, if you need the latest features, you can install it from the repository.
```
pip install git+https://github.com/probabilists/zuko
```
## Getting started
Normalizing flows are provided in the `zuko.flows` module. To build one, supply the number of sample and context features as well as the transformations' hyperparameters. Then, feeding a context $c$ to the flow returns a conditional distribution $p(x | c)$ which can be evaluated and sampled from.
```python
import torch
import zuko
# Neural spline flow (NSF) with 3 sample features and 5 context features
flow = zuko.flows.NSF(3, 5, transforms=3, hidden_features=[128] * 3)
# Train to maximize the log-likelihood
optimizer = torch.optim.Adam(flow.parameters(), lr=1e-3)
for x, c in trainset:
loss = -flow(c).log_prob(x) # -log p(x | c)
loss = loss.mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Sample 64 points x ~ p(x | c*)
x = flow(c_star).sample((64,))
```
Alternatively, flows can be built as custom `Flow` objects.
```python
from zuko.flows import Flow, MaskedAutoregressiveTransform, Unconditional
from zuko.distributions import DiagNormal
from zuko.transforms import RotationTransform
flow = Flow(
transform=[
MaskedAutoregressiveTransform(3, 5, hidden_features=(64, 64)),
Unconditional(RotationTransform, torch.randn(3, 3)),
MaskedAutoregressiveTransform(3, 5, hidden_features=(64, 64)),
],
base=Unconditional(
DiagNormal,
torch.zeros(3),
torch.ones(3),
buffer=True,
),
)
```
For more information, check out the documentation and tutorials at [zuko.readthedocs.io](https://zuko.readthedocs.io).
### Available flows
| Class | Year | Reference |
|:-------:|:----:|-----------|
| `GMM` | - | [Gaussian Mixture Model](https://wikipedia.org/wiki/Mixture_model#Gaussian_mixture_model) |
| `NICE` | 2014 | [Non-linear Independent Components Estimation](https://arxiv.org/abs/1410.8516) |
| `MAF` | 2017 | [Masked Autoregressive Flow for Density Estimation](https://arxiv.org/abs/1705.07057) |
| `NSF` | 2019 | [Neural Spline Flows](https://arxiv.org/abs/1906.04032) |
| `NCSF` | 2020 | [Normalizing Flows on Tori and Spheres](https://arxiv.org/abs/2002.02428) |
| `SOSPF` | 2019 | [Sum-of-Squares Polynomial Flow](https://arxiv.org/abs/1905.02325) |
| `NAF` | 2018 | [Neural Autoregressive Flows](https://arxiv.org/abs/1804.00779) |
| `UNAF` | 2019 | [Unconstrained Monotonic Neural Networks](https://arxiv.org/abs/1908.05164) |
| `CNF` | 2018 | [Neural Ordinary Differential Equations](https://arxiv.org/abs/1806.07366) |
| `GF` | 2020 | [Gaussianization Flows](https://arxiv.org/abs/2003.01941) |
| `BPF` | 2020 | [Bernstein-Polynomial Normalizing Flows](https://arxiv.org/abs/2204.13939) |
## Contributing
If you have a question, an issue or would like to contribute, please read our [contributing guidelines](https://github.com/probabilists/zuko/blob/master/CONTRIBUTING.md).
Raw data
{
"_id": null,
"home_page": "",
"name": "zuko",
"maintainer": "",
"docs_url": null,
"requires_python": ">=3.8",
"maintainer_email": "",
"keywords": "torch,normalizing flows,probability,density,generative,deep learning",
"author": "The Probabilists",
"author_email": "theprobabilists@gmail.com",
"download_url": "https://files.pythonhosted.org/packages/d3/23/d27153d5d4b67551e6541e1dd1d59101cd504389b4256de5340c94c8cd4e/zuko-1.1.0.tar.gz",
"platform": null,
"description": "\n\n# Zuko - Normalizing flows in PyTorch\n\nZuko is a Python package that implements normalizing flows in [PyTorch](https://pytorch.org). It relies as much as possible on distributions and transformations already provided by PyTorch. Unfortunately, the `Distribution` and `Transform` classes of `torch` are not sub-classes of `torch.nn.Module`, which means you cannot send their internal tensors to GPU with `.to('cuda')` or retrieve their parameters with `.parameters()`. Worse, the concepts of conditional distribution and transformation, which are essential for probabilistic inference, are impossible to express.\n\nTo solve these problems, `zuko` defines two concepts: the `LazyDistribution` and `LazyTransform`, which are any modules whose forward pass returns a `Distribution` or `Transform`, respectively. Because the creation of the actual distribution/transformation is delayed, an eventual condition can be easily taken into account. This design enables lazy distributions, including normalizing flows, to act like distributions while retaining features inherent to modules, such as trainable parameters. It also makes the implementations easy to understand and extend.\n\n> In the [Avatar](https://wikipedia.org/wiki/Avatar:_The_Last_Airbender) cartoon, [Zuko](https://wikipedia.org/wiki/Zuko) is a powerful firebender \ud83d\udd25\n\n## Acknowledgements\n\nZuko takes significant inspiration from [nflows](https://github.com/bayesiains/nflows) and [Stefan Webb](https://github.com/stefanwebb)'s work in [Pyro](https://github.com/pyro-ppl/pyro) and [FlowTorch](https://github.com/facebookincubator/flowtorch).\n\n## Installation\n\nThe `zuko` package is available on [PyPI](https://pypi.org/project/zuko), which means it is installable via `pip`.\n\n```\npip install zuko\n```\n\nAlternatively, if you need the latest features, you can install it from the repository.\n\n```\npip install git+https://github.com/probabilists/zuko\n```\n\n## Getting started\n\nNormalizing flows are provided in the `zuko.flows` module. To build one, supply the number of sample and context features as well as the transformations' hyperparameters. Then, feeding a context $c$ to the flow returns a conditional distribution $p(x | c)$ which can be evaluated and sampled from.\n\n```python\nimport torch\nimport zuko\n\n# Neural spline flow (NSF) with 3 sample features and 5 context features\nflow = zuko.flows.NSF(3, 5, transforms=3, hidden_features=[128] * 3)\n\n# Train to maximize the log-likelihood\noptimizer = torch.optim.Adam(flow.parameters(), lr=1e-3)\n\nfor x, c in trainset:\n loss = -flow(c).log_prob(x) # -log p(x | c)\n loss = loss.mean()\n\n optimizer.zero_grad()\n loss.backward()\n optimizer.step()\n\n# Sample 64 points x ~ p(x | c*)\nx = flow(c_star).sample((64,))\n```\n\nAlternatively, flows can be built as custom `Flow` objects.\n\n```python\nfrom zuko.flows import Flow, MaskedAutoregressiveTransform, Unconditional\nfrom zuko.distributions import DiagNormal\nfrom zuko.transforms import RotationTransform\n\nflow = Flow(\n transform=[\n MaskedAutoregressiveTransform(3, 5, hidden_features=(64, 64)),\n Unconditional(RotationTransform, torch.randn(3, 3)),\n MaskedAutoregressiveTransform(3, 5, hidden_features=(64, 64)),\n ],\n base=Unconditional(\n DiagNormal,\n torch.zeros(3),\n torch.ones(3),\n buffer=True,\n ),\n)\n```\n\nFor more information, check out the documentation and tutorials at [zuko.readthedocs.io](https://zuko.readthedocs.io).\n\n### Available flows\n\n| Class | Year | Reference |\n|:-------:|:----:|-----------|\n| `GMM` | - | [Gaussian Mixture Model](https://wikipedia.org/wiki/Mixture_model#Gaussian_mixture_model) |\n| `NICE` | 2014 | [Non-linear Independent Components Estimation](https://arxiv.org/abs/1410.8516) |\n| `MAF` | 2017 | [Masked Autoregressive Flow for Density Estimation](https://arxiv.org/abs/1705.07057) |\n| `NSF` | 2019 | [Neural Spline Flows](https://arxiv.org/abs/1906.04032) |\n| `NCSF` | 2020 | [Normalizing Flows on Tori and Spheres](https://arxiv.org/abs/2002.02428) |\n| `SOSPF` | 2019 | [Sum-of-Squares Polynomial Flow](https://arxiv.org/abs/1905.02325) |\n| `NAF` | 2018 | [Neural Autoregressive Flows](https://arxiv.org/abs/1804.00779) |\n| `UNAF` | 2019 | [Unconstrained Monotonic Neural Networks](https://arxiv.org/abs/1908.05164) |\n| `CNF` | 2018 | [Neural Ordinary Differential Equations](https://arxiv.org/abs/1806.07366) |\n| `GF` | 2020 | [Gaussianization Flows](https://arxiv.org/abs/2003.01941) |\n| `BPF` | 2020 | [Bernstein-Polynomial Normalizing Flows](https://arxiv.org/abs/2204.13939) |\n\n## Contributing\n\nIf you have a question, an issue or would like to contribute, please read our [contributing guidelines](https://github.com/probabilists/zuko/blob/master/CONTRIBUTING.md).\n",
"bugtrack_url": null,
"license": "",
"summary": "Normalizing flows in PyTorch",
"version": "1.1.0",
"project_urls": {
"Documentation": "https://zuko.readthedocs.io",
"Source": "https://github.com/probabilists/zuko",
"Tracker": "https://github.com/probabilists/zuko/issues"
},
"split_keywords": [
"torch",
"normalizing flows",
"probability",
"density",
"generative",
"deep learning"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "9e8d043b41ebf976a18a6d2d8a8042c8788da99121423ace642a3a7efa20789d",
"md5": "0b66cbf84491f095e9d5273bbf439fee",
"sha256": "7794969a8bb6e42ba3b721a128c62b3a2a63836b78db843b7e5defd0dfd572a8"
},
"downloads": -1,
"filename": "zuko-1.1.0-py3-none-any.whl",
"has_sig": false,
"md5_digest": "0b66cbf84491f095e9d5273bbf439fee",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.8",
"size": 38051,
"upload_time": "2024-01-26T10:42:07",
"upload_time_iso_8601": "2024-01-26T10:42:07.537339Z",
"url": "https://files.pythonhosted.org/packages/9e/8d/043b41ebf976a18a6d2d8a8042c8788da99121423ace642a3a7efa20789d/zuko-1.1.0-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "d323d27153d5d4b67551e6541e1dd1d59101cd504389b4256de5340c94c8cd4e",
"md5": "9ac7453ae78b64ad6d622f9f5845c274",
"sha256": "cd282dbbb4c845d2a104070b7780f128420b27296fda5b1731f0c84dbc99388e"
},
"downloads": -1,
"filename": "zuko-1.1.0.tar.gz",
"has_sig": false,
"md5_digest": "9ac7453ae78b64ad6d622f9f5845c274",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.8",
"size": 34320,
"upload_time": "2024-01-26T10:42:09",
"upload_time_iso_8601": "2024-01-26T10:42:09.339123Z",
"url": "https://files.pythonhosted.org/packages/d3/23/d27153d5d4b67551e6541e1dd1d59101cd504389b4256de5340c94c8cd4e/zuko-1.1.0.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-01-26 10:42:09",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "probabilists",
"github_project": "zuko",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"requirements": [],
"lcname": "zuko"
}
JSON
