The following sections provide greater details of the dataset, running training and inference, and the training results.

Scripts and sample code

In the root directory, the most important files are:

• Dockerfile - definition of the Docker image with all dependencies needed to run MoFlow
• setup.py - script that allows installing MoFlow with pip. Note that it does not include dependencies.

The moflow directory contains the definition of the network and tools needed for using it

• config.py - configuration of the dataset and network
• data - directory with tools needed to process and load the data
• model - directory with the definition of the MoFlow's building blocks and helper functions
• runtime - directory that contains code for running experiments, multi-GPU training, and logging. The most important files in this directory are train.py and generate.py, which allow running training or inference, respectively.
• utils.py- various helper functions

The scripts directory contains scripts for running the most typical workflows inside the docker container:

• benchmark_inference.sh and benchmark_training.sh for measuring the performance of inference or training, respectively
• data_preprocess.py for dataset preparation
• prepare_datasets.sh for downloading and preprocessing the data (note, that it launches data_preprocess.py)
• train.sh for launching training
• predict.sh for sampling random molecules from the trained model

Parameters

The complete list of parameters accepted by the runtime scripts (moflow/runtime/train.py and moflow/runtime/generate.py) consists of:

• --data_dir - Location for the dataset.
• --config_name - The config to choose. This parameter allows one to switch between different datasets and their dedicated configurations of the neural network. By default, a pre-defined "zinc250k" config is used.
• --results_dir - Directory where checkpoints are stored.
• --predictions_path - Path to store generated molecules. If an empty string is provided, predictions will not be saved (useful for benchmarking and debugging).
• --log_path - Path for DLLogger log. This file will contain information about the speed and accuracy of the model during training and inference. Note that if the file already exists, new logs will be added at the end.
• --log_interval - Frequency for writing logs, expressed in steps.
• --warmup_steps - Number of warmup steps. This value is used for benchmarking and for CUDA graph capture.
• --steps - Number of steps used for training/inference. This parameter allows finishing training earlier than the specified number of epochs. If used with inference, it allows generating more molecules (by default only a single batch of molecules is generated).
• --save_epochs - Frequency for saving checkpoints, expressed in epochs. If -1 is provided, checkpoints will not be saved.
• --eval_epochs - Evaluation frequency, expressed in epochs. If -1 is provided, an evaluation will not be performed.
• --learning_rate - Base learning rate.
• --beta1 - beta1 parameter for the Adam optimizer.
• --beta2 - beta2 parameter for the Adam optimizer.
• --clip - Gradient clipping norm.
• --epochs - Number of training epochs. Note that you can finish training mid-epoch by using "--steps" flag.
• --batch_size - Batch size per GPU.
• --num_workers - Number of workers in the data loader.
• --seed - Random seed used to initialize the distributed loaders.
• --local_rank - rank of the GPU, used to launch distributed training. This argument is specified automatically by torchrun and does not have to be provided by the user.
• --temperature - Temperature used for sampling.
• --val_batch_size - Number of molecules to generate during the validation step.
• --allow_untrained - Allow sampling molecules from an untrained network. Useful for performance benchmarking or debugging purposes.
• --correct_validity - Apply validity correction after the generation of the molecules.
• --amp - Use Automatic Mixed Precision
• --cuda_graph - Capture GPU kernels with CUDA graphs. This option allows to speed up training.
• --jit - Compile the model with torch.jit.script. Can be used to speed up training or inference.
• --verbosity - Verbosity level. Specify the following values: 0, 1, 2, 3, where 0 means minimal verbosity (errors only) and 3 - maximal (debugging).

Command-line options

To view the full list of available options and their descriptions, use the -h or --help command-line option, for example: python moflow/runtime/train.py --help

The following example output is printed when running the model:

usage: train.py [-h] [--data_dir DATA_DIR] [--config_name {zinc250k}] [--results_dir RESULTS_DIR] [--predictions_path PREDICTIONS_PATH] [--log_path LOG_PATH] [--log_interval LOG_INTERVAL]
                [--warmup_steps WARMUP_STEPS] [--steps STEPS] [--save_epochs SAVE_EPOCHS] [--eval_epochs EVAL_EPOCHS] [--learning_rate LEARNING_RATE] [--beta1 BETA1] [--beta2 BETA2] [--clip CLIP]
                [--epochs EPOCHS] [--batch_size BATCH_SIZE] [--num_workers NUM_WORKERS] [--seed SEED] [--local_rank LOCAL_RANK] [--temperature TEMPERATURE] [--val_batch_size VAL_BATCH_SIZE]
                [--allow_untrained] [--correct_validity] [--amp] [--cuda_graph] [--jit] [--verbosity {0,1,2,3}]

optional arguments:
  -h, --help            show this help message and exit
  --data_dir DATA_DIR   Location for the dataset.
  --config_name {zinc250k}
                        The config to choose. This parameter allows one to switch between different datasets and their dedicated configurations of the neural network. By default, a pre-defined
                        "zinc250k" config is used.
  --results_dir RESULTS_DIR
                        Directory where checkpoints are stored.
  --predictions_path PREDICTIONS_PATH
                        Path to store generated molecules. If an empty string is provided, predictions will not be saved (useful for benchmarking and debugging).
  --log_path LOG_PATH   Path for DLLogger log. This file will contain information about the speed and accuracy of the model during training and inference. Note that if the file already exists, new logs
                        will be added at the end.
  --log_interval LOG_INTERVAL
                        Frequency for writing logs, expressed in steps.
  --warmup_steps WARMUP_STEPS
                        Number of warmup steps. This value is used for benchmarking and for CUDA graph capture.
  --steps STEPS         Number of steps used for training/inference. This parameter allows finishing training earlier than the specified number of epochs. If used with inference, it allows generating
                        more molecules (by default only a single batch of molecules is generated).
  --save_epochs SAVE_EPOCHS
                        Frequency for saving checkpoints, expressed in epochs. If -1 is provided, checkpoints will not be saved.
  --eval_epochs EVAL_EPOCHS
                        Evaluation frequency, expressed in epochs. If -1 is provided, an evaluation will not be performed.
  --learning_rate LEARNING_RATE
                        Base learning rate.
  --beta1 BETA1         beta1 parameter for the optimizer.
  --beta2 BETA2         beta2 parameter for the optimizer.
  --clip CLIP           Gradient clipping norm.
  --epochs EPOCHS       Number of training epochs. Note that you can finish training mid-epoch by using "--steps" flag.
  --batch_size BATCH_SIZE
                        Batch size per GPU.
  --num_workers NUM_WORKERS
                        Number of workers in the data loader.
  --seed SEED           Random seed used to initialize the distributed loaders.
  --local_rank LOCAL_RANK
                        rank of the GPU, used to launch distributed training. This argument is specified automatically by `torchrun` and does not have to be provided by the user.
  --temperature TEMPERATURE
                        Temperature used for sampling.
  --val_batch_size VAL_BATCH_SIZE
                        Number of molecules to generate during validation step.
  --allow_untrained     Allow sampling molecules from an untrained network. Useful for performance benchmarking or debugging purposes.
  --correct_validity    Apply validity correction after the generation of the molecules.
  --amp                 Use Automatic Mixed Precision.
  --cuda_graph          Capture GPU kernels with CUDA graphs. This option allows to speed up training.
  --jit                 Compile the model with `torch.jit.script`. Can be used to speed up training or inference.
  --verbosity {0,1,2,3}
                        Verbosity level. Specify the following values: 0, 1, 2, 3, where 0 means minimal verbosity (errors only) and 3 - maximal (debugging).

Getting the data

The MoFlow model was trained on the ZINC 250k dataset. The original data split was used, with 224569 molecules in the training set and 24887 molecules in the test set.

This repository contains the prepare_datasets.sh script that will automatically download and process the dataset. By default, data will be downloaded to the /data/ directory.

Dataset guidelines

The dataset preparation is implemented in the scripts/data_preprocess.py script, and the parameters for the dataset are defined in the moflow/config.py file. The config includes information about data location, the structure of the CSV file, types and numbers of atoms in the molecules, and the number of nodes in the output graphs.

Initially, the data is stored in a CSV file that contains the molecules in SMILES format, together with their properties (optional). The data is loaded using the pandas library, and the SMILES strings are converted to molecules with RDKit.

Then, the molecules are converted into graphs with features assigned to nodes and edges. The first step is the standardization of molecular structures - each molecule is converted into canonical SMILES and loaded back, and kekulized. Then, two numpy arrays are constructed. The first array is a vector corresponding to graph nodes and contains atomic numbers for all atoms in the molecule. The second array is a 2D square matrix corresponding to graph edges and contains codes for atomic bond orders - 0 if two atoms are not connected, 1 for a single bond, 2 for a double bond, and 3 for a triple bond.

Both arrays are padded to some predefined size larger than the maximum number of atoms in the molecules in the dataset. For ZINC 250k, the maximum number of atoms is 38, and the output size of the numpy arrays is set to 40 for the nodes array and 40x40 for the edges array.

This representation of the data is dumped on the disk using the numpy savez function.

During training, the numpy arrays are loaded, and one-hot-encoding is used to represent atomic numbers (node features) and bond orders (edge features). This representation is then used for training the neural network.

Training process

The training script is located in moflow/runtime/train.py and it accepts the parameters listed above.

To make the usage of the model easier, there is also scripts/train.sh script that runs training with the default configuration and the evaluation using the trained checkpoint at the end. This script can be run without any arguments - then it launches training on a single GPU and performance optimizations enabled - automatic mixed precision (AMP) and CUDA graph capture.

./scripts/train.sh

It is also possible to pass the number of GPUs and precision ("amp" or "full") that should be used for training. For example, to launch training with eight GPUs and AMP, run:

./scripts/train.sh 8

and to launch four GPU training with full precision, run:

./scripts/train.sh 4 full

These two arguments can also be followed by extra flags that will be passed to training and evaluation commands. For example, to train on eight GPUs with AMP, batch size of 2048 per GPU and save logs in /results/dll.json, run:

./scripts/train.sh 8 amp --batch_size 2048 --log_path /results/dll.json

Alternatively, you can launch training with moflow/runtime/train.py. To run the model with multiple GPUs, run:

torchrun --nproc_per_node=<# GPUs> moflow/runtime/train.py <arguments>

To enable mixed precision training, add --amp. You can also optimize the performance further by adding --cuda_graph or --jit flags to enable CUDA graph capture or just-in-time compilation, respectively.

Logs

By default, logs are printed to the screen and not saved on disk. If you want to store the logs, pass --log_path flag to scripts/train.sh or moflow/runtime/train.py.

Checkpoints

By default, the training script saves checkpoints inside /results every five epochs. The location of the checkpoints directory can be modified with --results_dir flag and saving interval with --save_epochs flag (pass -1 if you do not want to save checkpoints). Up to five most recent checkpoints are kept while the older ones are removed.

Evaluation

The following metrics are used to evaluate the model:

• Validity - the percentage of predictions corresponding to the correct molecular graph.
• Uniqueness - the percentage of valid molecules that is unique.
• Novelty - the percentage of valid and unique molecules not present in the training set.
• N.U.V - the percentage of valid, unique, and novel molecules.

During training, a single batch of molecules is generated every couple of epochs to assess two metrics: validity and uniqueness, as they are quick to calculate and track the training progress.

By default, the validation batch size is set to 100 molecules per GPU, and evaluation happens every five epochs. This can be changed with --val_batch_size and --eval_epochs flags, respectively. To disable evaluation, pass --eval_epochs -1.

If you use scripts/train.sh, there is also a final evaluation of the model done on 100 batches of molecules. This larger sample is evaluated with all metrics described above, and we use N.U.V as the main metric.

Alternatively, you can trigger evaluation manually by running moflow/runtime/evaluate.py script. Make sure that you pass the same value for --results_dir for both training and evaluation scripts.

Inference process

Inference can be run by launching the moflow/runtime/generate.py or scripts/predict.sh script. The first one provides more flexibility and accepts the arguments listed above. The second script allows you to easily run the default configuration with performance optimization (--jit flag) and molecule validity correction (--correct_validity). To generate a single batch of molecules with AMP and batch size of 512, run:

./scripts/predict.sh

You can also provide batch size and precision to use for predictions. For example, to generate 1000 molecules with full precision, run:

./scripts/predict.sh 1000 full

The script also allows you to pass extra flags to the generation. For example, to generate 10 batches of 1000 each and save predictions inside /results/predictions.smi, run:

./scripts/predict.sh 1000 amp --steps 10 --predictions_path /results/predictions.smi