[JTVAE] Update for DGL 0.5 (#91)

* Update * Fix * Fix * Fix * Fix * Update * Update * Fix * Fix * Fix * Fix * Fix * Update * Fix * Fix * Fix * Fix * Fix * Fix * Fix * Fix * Update * Update * Fix * Update * Update * Fix * Fix * Fix * Update * Update * Update * Fix * Fix * Fix * Update * Update * Fix * Update * Update * Update * Update * Fix * Update * Fix * Update * Fix * Fix * Fix * Update * Update * Fix * Fix * Update * Fix * Fix
awslabs · Sep 14, 2020 · 7a910d0 · 7a910d0
1 parent b532efe
commit 7a910d0
Show file tree

Hide file tree

Showing 39 changed files with 2,022 additions and 2,214 deletions.
diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md
@@ -11,6 +11,6 @@ Contributors
 * [Jiajing Hu](https://github.com/jjhu94): Weave
 * [Zhaoqiang Chen](https://github.com/autodataming): Update RDkit channel for conda installation; Suggestions and improvement for rexgen_direct
 * [Yue Zhong](https://github.com/YueZhong-bio): HadamardLinkPredictor, link prediction for ogbl-ppa
-* [Yangkang Zhang](https://github.com/yangkang98): GNNOGBPredictor, graph property prediction for ogbg-ppa, support for GuacaMol in JTNN
+* [Yangkang Zhang](https://github.com/yangkang98): GNNOGBPredictor, graph property prediction for ogbg-ppa, support for GuacaMol in JTVAE
 * [Sooheon Kim](https://github.com/sooheon): Parallel processing for dataset construction from a CSV file
 * [Krishna Sirumalla](https://github.com/skrsna): Fix for regression_inference.py and classification_inference.py; support non-ring systems for ScaffoldSplitter
diff --git a/README.md b/README.md
@@ -28,7 +28,7 @@ DGL-LifeSci should work on
 * macOS X
 * Windows 10
 
-DGL-LifeSci requires python 3.6+, DGL 0.5.0+ and PyTorch 1.5.0+.
+DGL-LifeSci requires python 3.6+, DGL 0.5.2+ and PyTorch 1.5.0+.
 
 [Install pytorch](https://pytorch.org/get-started/locally/)
 
@@ -192,6 +192,6 @@ Below we provide some reference numbers to show how DGL improves the speed of tr
 | ---------------------------------- | ----------------------- | -------------------------- | ---------------------------- |
 | GCN on Tox21                       | 5.5 (DeepChem)          | 1.0                        | 5.5x                         |
 | AttentiveFP on Aromaticity         | 6.0                     | 1.2                        | 5x                           |
-| JTNN on ZINC                       | 1826                    | 743                        | 2.5x                         |
+| JTVAE on ZINC                      | 1826                    | 743                        | 2.5x                         |
 | WLN for reaction center prediction | 11657                   | 858 (1 GPU) / 134 (8 GPUs) | 13.6x (1GPU) / 87.0x (8GPUs) |
 | WLN for candidate ranking          | 40122                   | 22268                      | 1.8x                         |
diff --git a/docs/source/api/data.rst b/docs/source/api/data.rst
@@ -148,6 +148,21 @@ Adapting to New Datasets for Weisfeiler-Lehman Networks
 .. autoclass:: dgllife.data.WLNRankDataset
     :members: ignore_large, __getitem__, __len__
 
+Generative Models
+-----------------
+
+JTVAE
+`````
+
+.. autoclass:: dgllife.data.DGLMolTree
+    :members: treesize, recover, assemble
+
+.. autoclass:: dgllife.data.JTVAEDataset
+    :members: __len__, __getitem__, move_to_device
+
+.. autoclass:: dgllife.data.JTVAECollator
+    :members: __call__
+
 Protein-Ligand Binding Affinity Prediction
 ------------------------------------------
 

diff --git a/examples/README.md b/examples/README.md
@@ -57,7 +57,7 @@ We provide various examples across 3 applications -- property prediction, genera
     - [Example Training Script](generative_models/dgmg)
 - Junction Tree Variational Autoencoder for Molecular Graph Generation (JTNN) [[paper]](https://arxiv.org/abs/1802.04364)
     - [JTNN with DGL](../python/dgllife/model/model_zoo/jtnn)
-    - [Example Training Script](generative_models/jtnn)
+    - [Example Training Script](generative_models/jtvae)
 
 ## Binding Affinity Prediction
 

diff --git a/examples/generative_models/jtnn/jtnn/__init__.py b/examples/generative_models/jtnn/jtnn/__init__.py