Google Summer of Code 2022: Porting of Normalizing Flow Model from Tensorflow to Pytorch

#1

About me:

Hello everybody!

My name is Jose Antonio Siguenza and I’ll be working during this summer as a contributor in DeepChem via Open Chemistry in GSoC 2022.

Currently, I’m cursing my last year studying Chemical Engineering in Ecuador. Also, I’ve been involved in research about topics related to data science, deep learning, and science at university. Over the past year, I’ve gotten to know DeepChem and I started contributing in December 2021. In general, this has been a fascinating learning path about DL models and their applications in life sciences, especially generative models. DeepChem’s community and documentation have made this journey more enjoyable.

Project Description:

Lately, DeepChem has decided to port some TensorFlow models to PyTorch. This project aims to successfully migrate one of the organization’s porting lists. This model will perform a mix of invertible transformations between the base and target distribution. In order to optimize the results, a Normalizing Flow object initialized with the flow layers and base distribution would be iterated computing the loss for each epoch. For this reason, transformations and flow layers will be created with their respective unittests. Finally, the implementation will include the respective documentation and a tutorial, if applicable, following DeepChem’s API.

Contact:
GitHub:@JoseAntonioSiguenza
E-mail : jasiguen@espol.edu.ec - jasiguenespol@gmail.com

2 Likes
#2

June 1, 2022: PR #2918 on NormalizingFlow and Affine classes

First PR which includes the NormalizingFlow model and Affine classes. This last one is a layer (bijective transformation) commonly used. Affine transformation is based on a geometric operation expressed as y = exp(a)*x + b, where a and b are known as scale and shift parameters, respectively. The logarithm of the determinant jacobian matrix is computed. This term is important to consider in the optimization loop. There are 11 failing checks about format and tests.

Given Feedback:

  • Split the PR into smaller ones
  • Add docstrings, type annotations, and tests
  • Add new layers/models to the docs folder
#3

June 10, 2022: Splitted initial PR and construct Affine class

Split NormalizingFlow and NormalizingFlowModel classes and focus on the Affine class.

Hands-on:

GSoC '22 Final Report: The Porting of Normalizing Flow Model from Tensorflow to Pytorch
#4

June 17, 2022: PR #2944 on NormalizingFlow model

Got merged in #2918, which is a bijective transformation (Affine transformation) used in the Normalizing Flow model as a layer. Also, I’ve opened #2944, this PR mainly contains documentation and a draft of a NormalizingFlow model, type annotations, docstrings, and code conventions.

To do:

  • Add unittests
  • Add doctests
  • Improve docstrings (too general)
  • Update docs folder (models.rst)
1 Like
#5

June 24, 2022: PR #2944 on NormalizingFlow model

This week, the successful completion of the model was achieved. This is to compose the transformation of the layer as an nn.ModuleList class. Then, the created NormalizingFlow model performs two methods:

  • Sampling:

    • Receives as a parameter the number of samples (int) to compute.
    • Returns a tuple (samples (n_samples, dim), log_prob(n_samples)). The sample object is a torch.Tensor of N points after transforming the base_distribution according to layers. Log_prob is the logarithm of the jacobian determinant (deviation from between the base and transformed distributions)

  • Log_prob:

    • Receives a tensor as input. This tensor will be evaluated with the inverse pass (this is very important for the optimization loop when given a target distribution)

  • Documentation and test

    • Add an example of usage performing an Affine Transformation.
    • Update format according to numpy style.
    • Type annotations

  • To do:

    • Improve the test by changing the assert condition.
    • Include Transform class to perform forward and inverse passes the transformation layers.
    • Add type hints to base_distribution parameter
1 Like
#6

July 1, 2022: PR #2944 Merged In
The test on the Normalizing Flow model passed and finally got merged in. This test evaluates:

  • Sampling is achieved correctly according to the size of inputs/outputs and the log_prob evaluation (forward pass).
  • Log_prob of the inverse pass.
    • This result must be a tensor full of zeros because the optimization loop is not being applied.
  • To do:
    • Future PR
      • Include Transform class to perform forward and inverse passes for each transformation layer type.
      • Add type hints to base_distribution parameter.
    • Jupyter notebook tutorial comparing results with existing Tensorflow NormalizingFlow model.
#7

July 8, 2022: Design the Optimization Loop and RNVP Transformation

Start working with NormalizingFlow tutorial approaching with Pytorch.

  • The used optimization loop needs to be changed since it is based on Tensorflow.
  • Manage to work with NumpyDataset DeepChem object.

Design an RNVP transformation class with its methods used by the NormalizingFlow model.

  • Log_prob() and sample(), this methods should look similar to the Affine class methods.

To do:

  • Open a PR where RNVP class is defined with documentation and test. This PR should be placed at layers.py path.

Use the NormalizingFlow model and its transformation to have a look at the first results

1 Like
#8

July 15, 2022: Open #2996 (RNVP Transformation Layer)

Design an RNVP transformation class with its methods used by the NormalizingFlow model. This class is being modified at layers path. Real NVP transformations give advantages over conventional layers because of the ease to perform the inverse pass.

  • forward() and inverse() methods accomplishes the transformations.

  • Tests must be written

  • To do:

    • Once the layer is functional, a RealNVP model must be crated (this is a type of normalizing flow)
#9

July 22, 2022: #2996 Complete RNVP Transformation

Successfully written RNVP transformation class with its methods used by the NormalizingFlow model. When composing this layer, imagine it as a mask and each layer of the mask requires input. This class is being modified at layers.py path.

  • Log_det_jac must be to zero since nothing is being learned yet, but the test is not passing because some values are near zero (e-36).

  • To do:

    • RealNVP model (PR)
#10

July 29, 2022: #2996 Complete Real NVP Transformation

Successfully written RNVP transformation class with its methods, documentation and tests used by the NormalizingFlow model. The next step is to code the Real NVP algorithm which will perform the following methods:

  • sample(): Performs a sampling of the transformed distribution

  • Log_prob: Computes the log_prob of each sample. This is needed for the optimization loop.

  • To do:

    • RealNVP model (Open a PR)
    • Tutorial based on the existing NormalizingFlow tutorial using Pytorch.
#11

August 05, 2022: Complete PR for Real NVP model

Write the code for the Real NVP Normalizing Flow model with its methods:

  • sampling(): Performs a sampling of the transformed distribution and computes the transformed input using the forward pass and its determinant jacobian.

  • log_proba(): Computes the log_prob of each sample. This is needed for the optimization loop.

  • sample_for_each_step(): Performs a sampling of the transformed input for each layer of the mask.

  • To do:

    • RealNVP model (Open a PR)
    • Add test for this class.
#12

August 19, 2022: Updated PR #2996 of Real NVP layer

Fixes to PR #2996 in documentation and doctest has been applied.

  • The parameters for the forward and inverse pass of RealNVP Layer should be a tensor
  • This must return a transformed tensor with the computation of log_det_jacobian.
  • Doctest of this RealNVP layer has been changed with an example of its use.
  • Test needs to be modified for the flaky test to pass.
  • To do:
    • RealNVP model (Open a PR)
    • Add test for this class.
#13

August 26, 2022: PR #2996 Got Merged In

Finally got merged #2996, where RealNVP Layer and also minor fixes to NormalizingFlow model were applied.

  • This layer give advantages over non linear transformations in terms of calculation time and density estimation.
  • Currently, working in the tutorial and have the model running for a toy example (i. e. Moon Dataset)
  • I’ve been working on Google Colab with a ‘!pip install - - pre deepchem’ installation.
  • Evaluate the results on the QM9 dataset as N.C. Frey in his tutorial.
  • To do:
    • Complete successfully the tutorial of ‘Normalizing Flows on Molecular Data using Pytorch’
1 Like
#14

September 2, 2022: Tutorial Draft of the Normalizing Flow Model using Pytorch on Molecular Data

As the model and its layer were merged, the next step is to evaluate the performance on QM9 dataset. A tutorial was created to see the results, and a primary draft has a successfully running.

  • Some problems due to Pytorch “NaN” were detected.
  • Splitting the data in small batches to train the model solved this problem, but it also considers small molecules (needs to be polished).
  • To do:
    • Get to run the tutorial to predict bigger molecules.
#15

September 9, 2022: Tutorial Draft of the Normalizing Flow Model using Pytorch on Molecular Data

As the model and its layer were merged, the next step is to evaluate the performance on QM9 dataset. A jupyter notebook tutorial was created to see the results, and it is successfully running. Work must be done on this tutorial before merging into the main.

  • Splitting the data into small batches to train the model solved this problem, but it also considers small molecules.
  • To do:
    • Add transformation layers to the main repo, so the model can predict bigger molecules.
    • Minor fixes on the Normalizing Flow model.