Data

Tools

Indexes

Applications

Experiments

Awesome

An open API service indexing awesome lists of open source software.

https://github.com/msmrexe/pytorch-normalizing-flows-2d

A from-scratch implementation of a Normalizing Flow (Real NVP) to model and sample from a 2D 'two-moons' distribution. Allows for GIF visualization of the NF evolution.
https://github.com/msmrexe/pytorch-normalizing-flows-2d

change-of-variables course-project density-estimation generative-models maximum-likelihood-estimation normalizing-flows pytorch real-nvp university-project

Last synced: 3 months ago
JSON representation

A from-scratch implementation of a Normalizing Flow (Real NVP) to model and sample from a 2D 'two-moons' distribution. Allows for GIF visualization of the NF evolution.

Awesome Lists containing this project

README 

          
# Normalizing Flow for 2D Density Estimation



A from-scratch implementation of a Normalizing Flow model, based on Real NVP (Real Non-Volume Preserving) transformations, to learn and sample from a 2D 'two-moons' distribution. This project was developed for the M.S. Generative Models course at Sharif University of Technology.



## Features



* **Real NVP Implementation:** Core logic built using affine coupling layers.

* **Modular `Flow` Class:** A `nn.Module` that stacks multiple transformations for complex density estimation.

* **Maximum Likelihood Training:** Optimized by minimizing the negative log-likelihood (NLL) of the data.

* **Density Visualization:** Generates and saves the learned probability density at each epoch.

* **GIF Generation:** Automatically creates a GIF visualizing the evolution of the learned density over training.

* **Data Sampling:** Capable of sampling new data points from the learned distribution.



## Core Concepts & Techniques



* **Normalizing Flows:** A class of generative models that learn a data distribution $p_X(x)$ by transforming a simple base distribution $p_Z(z)$ (e.g., a Gaussian) through an invertible, differentiable mapping $f: \mathcal{Z} \to \mathcal{X}$.

* **Change of Variables Formula:** The core principle used for training. The log-likelihood of a data point $x$ is given by the formula:



  $$\log p_X(x) = \log p_Z(f(x)) + \log \left| \det J_f(x) \right|$$



  where $z = f(x)$ is the forward transformation (data to latent) and $J_f(x)$ is its Jacobian.

  

* **Coupling Layers (Real NVP):** An efficient invertible transformation $y = f(x)$ where the input $x$ is split into two parts ($x_A, x_B$). One part is left unchanged ($y_A = x_A$), while the other is transformed using a simple affine function (scaling and translation) whose parameters ($s, t$) are generated by a neural network that only takes the unchanged part as input:



  $$y_B = x_B \odot \exp(s(x_A)) + t(x_A)$$



  This design results in a triangular Jacobian matrix, making the determinant (and its log) trivial to compute: it is simply the sum of the scaling factors, $\sum s(x_A)$.



* **Maximum Likelihood Estimation (MLE):** The model is trained by adjusting the parameters of the $s$ and $t$ networks to maximize the log-likelihood of the training data, which is equivalent to minimizing the negative log-likelihood (NLL) loss.



---



## How It Works



This project learns a complex distribution (two moons) by mapping it to a simple one (a 2D standard Gaussian).



### 1. Core Logic & Architecture



The repository implements a `Flow` model composed of a series of `CouplingTransform` layers. Each `CouplingTransform` layer splits the 2D input $x = (x_1, x_2)$ based on an alternating binary mask.



1.  **Transform 1 (Mask `[1, 0]`):**

    * $x_1$ is **fixed**.

    * $s_1, t_1 = \text{NN}_1(x_1)$

    * $z_1 = x_1$

    * $z_2 = x_2 \odot \exp(s_1) + t_1$



2.  **Transform 2 (Mask `[0, 1]`):**

    * $z_2$ is **fixed**.

    * $s_2, t_2 = \text{NN}_2(z_2)$

    * $z'_1 = z_1 \odot \exp(s_2) + t_2$

    * $z'_2 = z_2$



This sequence of invertible transformations ($x \to z \to z' \dots$) maps the complex 'two-moons' data to a 2D standard Gaussian. Training is done by maximizing $\log p_X(x)$, and sampling is done by drawing $z$ from the Gaussian and applying the inverse of all transforms: $x = f_1^{-1}(f_2^{-1}(\dots f_N^{-1}(z)))$.



### 2. Key Files & Functions



* `main.py`: The main executable script. It handles argument parsing, sets up logging, loads the data, initializes the `Flow` model, and runs the training loop. After training, it generates all final visualizations.

* `src/model.py`: Defines the `Flow` class. This class stacks multiple `CouplingTransform` layers. It implements:

    * `forward(x)`: Maps data to latent space ($x \to z$) and returns $z$ and the total $\log \det J$.

    * `inverse(z)`: Maps latent space to data space ($z \to x$).

    * `log_prob(x)`: Applies the Change of Variables formula to compute the log-likelihood of the data.

    * `sample(n)`: Generates $n$ new samples by drawing from the base distribution and applying the `inverse` method.

* `src/transforms.py`: Defines the `CouplingTransform` class. This class implements the core Real NVP logic, including the $s$ and $t$ neural networks, the `forward` pass, the `inverse` pass, and the log-det-Jacobian calculation.

* `src/utils.py`: Contains all helper functions for loading the `make_moons` dataset (`load_data`), creating the alternating masks (`create_alternating_masks`), setting up logging (`setup_logging`), and handling all visualizations (`visualize_density`, `plot_comparison`, `create_gif`).



---



## Project Structure



```

pytorch-normalizing-flows-2d/

├── .gitignore             # Ignores Python caches, envs, and output files

├── LICENSE                # MIT License file

├── README.md              # This repository README

├── main.py                # Main script to train the model and generate outputs

├── requirements.txt       # Required Python packages

├── run_flow_model.ipynb   # Jupyter Notebook to run the training and view results

└── src/

    ├── __init__.py        # Makes 'src' a Python package

    ├── model.py           # Contains the main 'Flow' model class

    ├── transforms.py      # Defines the 'CouplingTransform' layer

    └── utils.py           # Helper functions for data, logging, and visualization

```



## How to Use



1.  **Clone the Repository:**

    ```bash

    git clone https://github.com/msmrexe/pytorch-normalizing-flows-2d.git

    cd pytorch-normalizing-flows-2d

    ```



2.  **Setup the Environment:**

    It's recommended to use a virtual environment.

    ```bash

    python -m venv venv

    source venv/bin/activate  # On Windows: venv\Scripts\activate

    pip install -r requirements.txt

    ```



3.  **Run the Training Script:**

    You can run the model directly from the command line.

    ```bash

    python main.py --epochs 150 --lr 0.001 --output_dir outputs

    ```

    * `--epochs`: Number of training epochs (e.g., 150).

    * `--lr`: Learning rate (e.g., 0.001).

    * `--output_dir`: Directory to save final plots and GIF.

    * `--skip_frames`: (Optional) Use this flag to disable saving frames every epoch, which speeds up training.



4.  **Example Output / Check Results:**

    The script will log progress to the console and to `logs/flow.log`. After training, you will find the following files in the `outputs/` directory:

    * `density_evolution.gif`: An animation of the learned density.

    * `data_comparison.png`: A scatter plot comparing original and generated data.

    * `final_density.png`: The final learned probability density.



    Alternatively, you can run the `run_flow_model.ipynb` notebook in Jupyter to execute the training and see the outputs displayed inline.



---



## Author



Feel free to connect or reach out if you have any questions!



* **Maryam Rezaee**

* **GitHub:** [@msmrexe](https://github.com/msmrexe)

* **Email:** [ms.maryamrezaee@gmail.com](mailto:ms.maryamrezaee@gmail.com)



---



## License



This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for full details.