Code is adapted from MarcusOlivecrona/REINVENT.
Updated for use with PyTorch v1.10.0. Made for use on CC cluster.
All required imports and modules are found in make_env.sh
. Create virtual environment using
./make_env.sh
source ~/env/reinvent/bin/activate
Put data into the data/
directory. Convert the smiles of data into data/data.smi
by running
python prepare_data.py
You may need to modify the script depending on the format of dataset.
Then run
python data_structs.py
to create the mols_filtered.smi
and Voc
file required in the data/
directory. Note that original filtering from REINVENT 1.0 is removed, and SMILES are just check for validitiy and canonicalised.
Pretrain the model for specified num-epochs
by running:
python train_prior.py --num-epochs 100
You can view progress bar by toggling flag --verbose
.
Early stopping is implemented, default split is 80/20 using the first 80% of SMILES as training. Learning rate is decreased every 50 steps. The model is saved in data/Prior.ckpt
.
The fitness function is a class in custom.py
. Modify the fitness function by changing the fitness_function
function, which takes in a SMILES and outputs a float value. Return None
for invalid smiles.
The values will be automatically scaled by a sigmoid. For sampled molecules m with fitness f, the score is (invalid molecules have 0 reward and do not affect the agent)
where a is the average of all known fitnesses from previous generations, and b is the slope of the sigmoid, calculated from a threshold.
For a set F of fitnesses from previous generations
where c is the specified threshold (default 0.8). This means for the current fitnesses F, the maximum fitness will map to 0.8 on the sigmoid. The smaller the threshold c, the stronger the reward for fitnesses larger than F.
All other parameters are left as default from REINVENT 1.0.
Run the REINVENT for generations = num-steps
with evaluations = batch_size
per generation. Run the script
python main.py --num-steps 10 --batch-size 500
Results will be saved in data/results
, in a folder named by the timestamp of run. In this folder, a trace of the optimisation is plotted, and the results for each generation (results.csv
), and the best molecule per generation (best_results.csv
) are saved.