CINXE.COM

<!DOCTYPE html> <html lang="en"> <head> <meta charset="utf-8"> <meta name="viewport" content="width=device-width, initial-scale=1"> <meta name="description" content="Keras documentation"> <meta name="author" content="Keras Team"> <link rel="shortcut icon" href="https://keras.io/img/favicon.ico"> <link rel="canonical" href="https://keras.io/examples/generative/molecule_generation/" />  <meta property="og:title" content="Keras documentation: Drug Molecule Generation with VAE"> <meta property="og:image" content="https://keras.io/img/logo-k-keras-wb.png"> <meta name="twitter:title" content="Keras documentation: Drug Molecule Generation with VAE"> <meta name="twitter:image" content="https://keras.io/img/k-keras-social.png"> <meta name="twitter:card" content="summary"> <title>Drug Molecule Generation with VAE</title>  <link href="/css/bootstrap.min.css" rel="stylesheet">  <link href="https://fonts.googleapis.com/css2?family=Open+Sans:wght@400;600;700;800&display=swap" rel="stylesheet">  <link href="/css/docs.css" rel="stylesheet"> <link href="/css/monokai.css" rel="stylesheet">  <script>(function(w,d,s,l,i){w[l]=w[l]||[];w[l].push({'gtm.start': new Date().getTime(),event:'gtm.js'});var f=d.getElementsByTagName(s)[0], j=d.createElement(s),dl=l!='dataLayer'?'&l='+l:'';j.async=true;j.src= 'https://www.googletagmanager.com/gtm.js?id='+i+dl;f.parentNode.insertBefore(j,f); })(window,document,'script','dataLayer','GTM-5DNGF4N'); </script> <script> (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){ (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o), m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m) })(window,document,'script','https://www.google-analytics.com/analytics.js','ga'); ga('create', 'UA-175165319-128', 'auto'); ga('send', 'pageview'); </script>  <script async defer src="https://buttons.github.io/buttons.js"></script> </head> <body>  <noscript><iframe src="https://www.googletagmanager.com/ns.html?id=GTM-5DNGF4N" height="0" width="0" style="display:none;visibility:hidden"></iframe></noscript>  <div class='k-page'> <div class="k-nav" id="nav-menu"> <a href='/'><img src='/img/logo-small.png' class='logo-small' /></a> <div class="nav flex-column nav-pills" role="tablist" aria-orientation="vertical"> <a class="nav-link" href="/about/" role="tab" aria-selected="">About Keras</a> <a class="nav-link" href="/getting_started/" role="tab" aria-selected="">Getting started</a> <a class="nav-link" href="/guides/" role="tab" aria-selected="">Developer guides</a> <a class="nav-link" href="/api/" role="tab" aria-selected="">Keras 3 API documentation</a> <a class="nav-link" href="/2.18/api/" role="tab" aria-selected="">Keras 2 API documentation</a> <a class="nav-link active" href="/examples/" role="tab" aria-selected="">Code examples</a> <a class="nav-sublink" href="/examples/vision/">Computer Vision</a> <a class="nav-sublink" href="/examples/nlp/">Natural Language Processing</a> <a class="nav-sublink" href="/examples/structured_data/">Structured Data</a> <a class="nav-sublink" href="/examples/timeseries/">Timeseries</a> <a class="nav-sublink active" href="/examples/generative/">Generative Deep Learning</a> <a class="nav-sublink2" href="/examples/generative/ddim/">Denoising Diffusion Implicit Models</a> <a class="nav-sublink2" href="/examples/generative/random_walks_with_stable_diffusion/">A walk through latent space with Stable Diffusion</a> <a class="nav-sublink2" href="/examples/generative/dreambooth/">DreamBooth</a> <a class="nav-sublink2" href="/examples/generative/ddpm/">Denoising Diffusion Probabilistic Models</a> <a class="nav-sublink2" href="/examples/generative/fine_tune_via_textual_inversion/">Teach StableDiffusion new concepts via Textual Inversion</a> <a class="nav-sublink2" href="/examples/generative/finetune_stable_diffusion/">Fine-tuning Stable Diffusion</a> <a class="nav-sublink2" href="/examples/generative/vae/">Variational AutoEncoder</a> <a class="nav-sublink2" href="/examples/generative/dcgan_overriding_train_step/">GAN overriding Model.train_step</a> <a class="nav-sublink2" href="/examples/generative/wgan_gp/">WGAN-GP overriding Model.train_step</a> <a class="nav-sublink2" href="/examples/generative/conditional_gan/">Conditional GAN</a> <a class="nav-sublink2" href="/examples/generative/cyclegan/">CycleGAN</a> <a class="nav-sublink2" href="/examples/generative/gan_ada/">Data-efficient GANs with Adaptive Discriminator Augmentation</a> <a class="nav-sublink2" href="/examples/generative/deep_dream/">Deep Dream</a> <a class="nav-sublink2" href="/examples/generative/gaugan/">GauGAN for conditional image generation</a> <a class="nav-sublink2" href="/examples/generative/pixelcnn/">PixelCNN</a> <a class="nav-sublink2" href="/examples/generative/stylegan/">Face image generation with StyleGAN</a> <a class="nav-sublink2" href="/examples/generative/vq_vae/">Vector-Quantized Variational Autoencoders</a> <a class="nav-sublink2" href="/examples/generative/neural_style_transfer/">Neural style transfer</a> <a class="nav-sublink2" href="/examples/generative/adain/">Neural Style Transfer with AdaIN</a> <a class="nav-sublink2" href="/examples/generative/gpt2_text_generation_with_keras_hub/">GPT2 Text Generation with KerasHub</a> <a class="nav-sublink2" href="/examples/generative/text_generation_gpt/">GPT text generation from scratch with KerasHub</a> <a class="nav-sublink2" href="/examples/generative/text_generation_with_miniature_gpt/">Text generation with a miniature GPT</a> <a class="nav-sublink2" href="/examples/generative/lstm_character_level_text_generation/">Character-level text generation with LSTM</a> <a class="nav-sublink2" href="/examples/generative/text_generation_fnet/">Text Generation using FNet</a> <a class="nav-sublink2 active" href="/examples/generative/molecule_generation/">Drug Molecule Generation with VAE</a> <a class="nav-sublink2" href="/examples/generative/wgan-graphs/">WGAN-GP with R-GCN for the generation of small molecular graphs</a> <a class="nav-sublink2" href="/examples/generative/random_walks_with_stable_diffusion_3/">A walk through latent space with Stable Diffusion 3</a> <a class="nav-sublink2" href="/examples/generative/real_nvp/">Density estimation using Real NVP</a> <a class="nav-sublink" href="/examples/audio/">Audio Data</a> <a class="nav-sublink" href="/examples/rl/">Reinforcement Learning</a> <a class="nav-sublink" href="/examples/graph/">Graph Data</a> <a class="nav-sublink" href="/examples/keras_recipes/">Quick Keras Recipes</a> <a class="nav-link" href="/keras_tuner/" role="tab" aria-selected="">KerasTuner: Hyperparameter Tuning</a> <a class="nav-link" href="/keras_hub/" role="tab" aria-selected="">KerasHub: Pretrained Models</a> <a class="nav-link" href="/keras_cv/" role="tab" aria-selected="">KerasCV: Computer Vision Workflows</a> <a class="nav-link" href="/keras_nlp/" role="tab" aria-selected="">KerasNLP: Natural Language Workflows</a> </div> </div> <div class='k-main'> <div class='k-main-top'> <script> function displayDropdownMenu() { e = document.getElementById("nav-menu"); if (e.style.display == "block") { e.style.display = "none"; } else { e.style.display = "block"; document.getElementById("dropdown-nav").style.display = "block"; } } function resetMobileUI() { if (window.innerWidth <= 840) { document.getElementById("nav-menu").style.display = "none"; document.getElementById("dropdown-nav").style.display = "block"; } else { document.getElementById("nav-menu").style.display = "block"; document.getElementById("dropdown-nav").style.display = "none"; } var navmenu = document.getElementById("nav-menu"); var menuheight = navmenu.clientHeight; var kmain = document.getElementById("k-main-id"); kmain.style.minHeight = (menuheight + 100) + 'px'; } window.onresize = resetMobileUI; window.addEventListener("load", (event) => { resetMobileUI() }); </script> <div id='dropdown-nav' onclick="displayDropdownMenu();"> <svg viewBox="-20 -20 120 120" width="60" height="60"> <rect width="100" height="20"></rect> <rect y="30" width="100" height="20"></rect> <rect y="60" width="100" height="20"></rect> </svg> </div> <form class="bd-search d-flex align-items-center k-search-form" id="search-form"> <input type="search" class="k-search-input" id="search-input" placeholder="Search Keras documentation..." aria-label="Search Keras documentation..." autocomplete="off"> <button class="k-search-btn"> <svg width="13" height="13" viewBox="0 0 13 13"><title>search</title><path d="m4.8495 7.8226c0.82666 0 1.5262-0.29146 2.0985-0.87438 0.57232-0.58292 0.86378-1.2877 0.87438-2.1144 0.010599-0.82666-0.28086-1.5262-0.87438-2.0985-0.59352-0.57232-1.293-0.86378-2.0985-0.87438-0.8055-0.010599-1.5103 0.28086-2.1144 0.87438-0.60414 0.59352-0.8956 1.293-0.87438 2.0985 0.021197 0.8055 0.31266 1.5103 0.87438 2.1144 0.56172 0.60414 1.2665 0.8956 2.1144 0.87438zm4.4695 0.2115 3.681 3.6819-1.259 1.284-3.6817-3.7 0.0019784-0.69479-0.090043-0.098846c-0.87973 0.76087-1.92 1.1413-3.1207 1.1413-1.3553 0-2.5025-0.46363-3.4417-1.3909s-1.4088-2.0686-1.4088-3.4239c0-1.3553 0.4696-2.4966 1.4088-3.4239 0.9392-0.92727 2.0864-1.3969 3.4417-1.4088 1.3553-0.011889 2.4906 0.45771 3.406 1.4088 0.9154 0.95107 1.379 2.0924 1.3909 3.4239 0 1.2126-0.38043 2.2588-1.1413 3.1385l0.098834 0.090049z"></path></svg> </button> </form> <script> var form = document.getElementById('search-form'); form.onsubmit = function(e) { e.preventDefault(); var query = document.getElementById('search-input').value; window.location.href = '/search.html?query=' + query; return False } </script> </div> <div class='k-main-inner' id='k-main-id'> <div class='k-location-slug'> ► <a href='/examples/'>Code examples</a> / <a href='/examples/generative/'>Generative Deep Learning</a> / Drug Molecule Generation with VAE </div> <div class='k-content'> <h1 id="drug-molecule-generation-with-vae">Drug Molecule Generation with VAE</h1> Author: <a href="https://www.linkedin.com/in/victor-basu-520958147">Victor Basu</a> Date created: 2022/03/10 Last modified: 2022/03/24 Description: Implementing a Convolutional Variational AutoEncoder (VAE) for Drug Discovery. <div class='example_version_banner keras_2'>ⓘ This example uses Keras 2</div> <img class="k-inline-icon" src="https://colab.research.google.com/img/colab_favicon.ico"/> <a href="https://colab.research.google.com/github/keras-team/keras-io/blob/master/examples/generative/ipynb/molecule_generation.ipynb">View in Colab</a> •<img class="k-inline-icon" src="https://github.com/favicon.ico"/> <a href="https://github.com/keras-team/keras-io/blob/master/examples/generative/molecule_generation.py">GitHub source</a> <hr /> <h2 id="introduction">Introduction</h2> In this example, we use a Variational Autoencoder to generate molecules for drug discovery. We use the research papers <a href="https://arxiv.org/abs/1610.02415">Automatic chemical design using a data-driven continuous representation of molecules</a> and <a href="https://arxiv.org/abs/1805.11973">MolGAN: An implicit generative model for small molecular graphs</a> as a reference. The model described in the paper Automatic chemical design using a data-driven continuous representation of molecules generates new molecules via efficient exploration of open-ended spaces of chemical compounds. The model consists of three components: Encoder, Decoder and Predictor. The Encoder converts the discrete representation of a molecule into a real-valued continuous vector, and the Decoder converts these continuous vectors back to discrete molecule representations. The Predictor estimates chemical properties from the latent continuous vector representation of the molecule. Continuous representations allow the use of gradient-based optimization to efficiently guide the search for optimized functional compounds. <img alt="intro" src="https://bit.ly/3CtPMzM" /> Figure (a) - A diagram of the autoencoder used for molecule design, including the joint property prediction model. Starting from a discrete molecule representation, such as a SMILES string, the encoder network converts each molecule into a vector in the latent space, which is effectively a continuous molecule representation. Given a point in the latent space, the decoder network produces a corresponding SMILES string. A multilayer perceptron network estimates the value of target properties associated with each molecule. Figure (b) - Gradient-based optimization in continuous latent space. After training a surrogate model <code>f(z)</code> to predict the properties of molecules based on their latent representation <code>z</code>, we can optimize <code>f(z)</code> with respect to <code>z</code> to find new latent representations expected to match specific desired properties. These new latent representations can then be decoded into SMILES strings, at which point their properties can be tested empirically. For an explanation and implementation of MolGAN, please refer to the Keras Example <a href="https://bit.ly/3pU6zXK">WGAN-GP with R-GCN for the generation of small molecular graphs</a> by Alexander Kensert. Many of the functions used in the present example are from the above Keras example. <hr /> <h2 id="setup">Setup</h2> RDKit is an open source toolkit for cheminformatics and machine learning. This toolkit come in handy if one is into drug discovery domain. In this example, RDKit is used to conveniently and efficiently transform SMILES to molecule objects, and then from those obtain sets of atoms and bonds. Quoting from <a href="https://keras.io/examples/generative/wgan-graphs/">WGAN-GP with R-GCN for the generation of small molecular graphs</a>): "SMILES expresses the structure of a given molecule in the form of an ASCII string. The SMILES string is a compact encoding which, for smaller molecules, is relatively human-readable. Encoding molecules as a string both alleviates and facilitates database and/or web searching of a given molecule. RDKit uses algorithms to accurately transform a given SMILES to a molecule object, which can then be used to compute a great number of molecular properties/features." <div class="codehilite"><pre><code>!pip -q install rdkit-pypi==2021.9.4 </code></pre></div> <div class="codehilite"><pre><code>import ast import pandas as pd import numpy as np import tensorflow as tf from tensorflow import keras from tensorflow.keras import layers import matplotlib.pyplot as plt from rdkit import Chem, RDLogger from rdkit.Chem import BondType from rdkit.Chem.Draw import MolsToGridImage RDLogger.DisableLog("rdApp.*") </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>[K |████████████████████████████████| 20.6 MB 1.2 MB/s [?25h </code></pre></div> </div> <hr /> <h2 id="dataset">Dataset</h2> We use the <a href="https://bit.ly/3IVBI4x">ZINC – A Free Database of Commercially Available Compounds for Virtual Screening</a> dataset. The dataset comes with molecule formula in SMILE representation along with their respective molecular properties such as logP (water–octanal partition coefficient), SAS (synthetic accessibility score) and QED (Qualitative Estimate of Drug-likeness). <div class="codehilite"><pre><code>csv_path = keras.utils.get_file( "/content/250k_rndm_zinc_drugs_clean_3.csv", "https://raw.githubusercontent.com/aspuru-guzik-group/chemical_vae/master/models/zinc_properties/250k_rndm_zinc_drugs_clean_3.csv", ) df = pd.read_csv("/content/250k_rndm_zinc_drugs_clean_3.csv") df["smiles"] = df["smiles"].apply(lambda s: s.replace("\n", "")) df.head() </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Downloading data from https://raw.githubusercontent.com/aspuru-guzik-group/chemical_vae/master/models/zinc_properties/250k_rndm_zinc_drugs_clean_3.csv 22606589/22606589 [==============================] - 0s 0us/step </code></pre></div> </div> <div> <style scoped> .dataframe tbody tr th:only-of-type { vertical-align: middle; } <div class="k-default-codeblock"> <div class="codehilite"><pre><code>.dataframe tbody tr th { vertical-align: top; } .dataframe thead th { text-align: right; } </code></pre></div> </div> </style> <table border="1" class="dataframe"> <thead> <tr style="text-align: right;"> <th></th> <th>smiles</th> <th>logP</th> <th>qed</th> <th>SAS</th> </tr> </thead> <tbody> <tr> <th>0</th> <td>CC(C)(C)c1ccc2occ(CC(=O)Nc3ccccc3F)c2c1</td> <td>5.05060</td> <td>0.702012</td> <td>2.084095</td> </tr> <tr> <th>1</th> <td>C[C@@H]1CC(Nc2cncc(-c3nncn3C)c2)C[C@@H](C)C1</td> <td>3.11370</td> <td>0.928975</td> <td>3.432004</td> </tr> <tr> <th>2</th> <td>N#Cc1ccc(-c2ccc(O[C@@H](C(=O)N3CCCC3)c3ccccc3)...</td> <td>4.96778</td> <td>0.599682</td> <td>2.470633</td> </tr> <tr> <th>3</th> <td>CCOC(=O)[C@@H]1CCCN(C(=O)c2nc(-c3ccc(C)cc3)n3c...</td> <td>4.00022</td> <td>0.690944</td> <td>2.822753</td> </tr> <tr> <th>4</th> <td>N#CC1=C(SCC(=O)Nc2cccc(Cl)c2)N=C([O-])[C@H](C#...</td> <td>3.60956</td> <td>0.789027</td> <td>4.035182</td> </tr> </tbody> </table> </div> <hr /> <h2 id="hyperparameters">Hyperparameters</h2> <div class="codehilite"><pre><code>SMILE_CHARSET = '["C", "B", "F", "I", "H", "O", "N", "S", "P", "Cl", "Br"]' bond_mapping = {"SINGLE": 0, "DOUBLE": 1, "TRIPLE": 2, "AROMATIC": 3} bond_mapping.update( {0: BondType.SINGLE, 1: BondType.DOUBLE, 2: BondType.TRIPLE, 3: BondType.AROMATIC} ) SMILE_CHARSET = ast.literal_eval(SMILE_CHARSET) MAX_MOLSIZE = max(df["smiles"].str.len()) SMILE_to_index = dict((c, i) for i, c in enumerate(SMILE_CHARSET)) index_to_SMILE = dict((i, c) for i, c in enumerate(SMILE_CHARSET)) atom_mapping = dict(SMILE_to_index) atom_mapping.update(index_to_SMILE) BATCH_SIZE = 100 EPOCHS = 10 VAE_LR = 5e-4 NUM_ATOMS = 120 # Maximum number of atoms ATOM_DIM = len(SMILE_CHARSET) # Number of atom types BOND_DIM = 4 + 1 # Number of bond types LATENT_DIM = 435 # Size of the latent space def smiles_to_graph(smiles): # Converts SMILES to molecule object molecule = Chem.MolFromSmiles(smiles) # Initialize adjacency and feature tensor adjacency = np.zeros((BOND_DIM, NUM_ATOMS, NUM_ATOMS), "float32") features = np.zeros((NUM_ATOMS, ATOM_DIM), "float32") # loop over each atom in molecule for atom in molecule.GetAtoms(): i = atom.GetIdx() atom_type = atom_mapping[atom.GetSymbol()] features[i] = np.eye(ATOM_DIM)[atom_type] # loop over one-hop neighbors for neighbor in atom.GetNeighbors(): j = neighbor.GetIdx() bond = molecule.GetBondBetweenAtoms(i, j) bond_type_idx = bond_mapping[bond.GetBondType().name] adjacency[bond_type_idx, [i, j], [j, i]] = 1 # Where no bond, add 1 to last channel (indicating "non-bond") # Notice: channels-first adjacency[-1, np.sum(adjacency, axis=0) == 0] = 1 # Where no atom, add 1 to last column (indicating "non-atom") features[np.where(np.sum(features, axis=1) == 0)[0], -1] = 1 return adjacency, features def graph_to_molecule(graph): # Unpack graph adjacency, features = graph # RWMol is a molecule object intended to be edited molecule = Chem.RWMol() # Remove "no atoms" & atoms with no bonds keep_idx = np.where( (np.argmax(features, axis=1) != ATOM_DIM - 1) & (np.sum(adjacency[:-1], axis=(0, 1)) != 0) )[0] features = features[keep_idx] adjacency = adjacency[:, keep_idx, :][:, :, keep_idx] # Add atoms to molecule for atom_type_idx in np.argmax(features, axis=1): atom = Chem.Atom(atom_mapping[atom_type_idx]) _ = molecule.AddAtom(atom) # Add bonds between atoms in molecule; based on the upper triangles # of the [symmetric] adjacency tensor (bonds_ij, atoms_i, atoms_j) = np.where(np.triu(adjacency) == 1) for (bond_ij, atom_i, atom_j) in zip(bonds_ij, atoms_i, atoms_j): if atom_i == atom_j or bond_ij == BOND_DIM - 1: continue bond_type = bond_mapping[bond_ij] molecule.AddBond(int(atom_i), int(atom_j), bond_type) # Sanitize the molecule; for more information on sanitization, see # https://www.rdkit.org/docs/RDKit_Book.html#molecular-sanitization flag = Chem.SanitizeMol(molecule, catchErrors=True) # Let's be strict. If sanitization fails, return None if flag != Chem.SanitizeFlags.SANITIZE_NONE: return None return molecule </code></pre></div> <hr /> <h2 id="generate-training-set">Generate training set</h2> <div class="codehilite"><pre><code>train_df = df.sample(frac=0.75, random_state=42) # random state is a seed value train_df.reset_index(drop=True, inplace=True) adjacency_tensor, feature_tensor, qed_tensor = [], [], [] for idx in range(8000): adjacency, features = smiles_to_graph(train_df.loc[idx]["smiles"]) qed = train_df.loc[idx]["qed"] adjacency_tensor.append(adjacency) feature_tensor.append(features) qed_tensor.append(qed) adjacency_tensor = np.array(adjacency_tensor) feature_tensor = np.array(feature_tensor) qed_tensor = np.array(qed_tensor) class RelationalGraphConvLayer(keras.layers.Layer): def __init__( self, units=128, activation="relu", use_bias=False, kernel_initializer="glorot_uniform", bias_initializer="zeros", kernel_regularizer=None, bias_regularizer=None, **kwargs ): super().__init__(**kwargs) self.units = units self.activation = keras.activations.get(activation) self.use_bias = use_bias self.kernel_initializer = keras.initializers.get(kernel_initializer) self.bias_initializer = keras.initializers.get(bias_initializer) self.kernel_regularizer = keras.regularizers.get(kernel_regularizer) self.bias_regularizer = keras.regularizers.get(bias_regularizer) def build(self, input_shape): bond_dim = input_shape[0][1] atom_dim = input_shape[1][2] self.kernel = self.add_weight( shape=(bond_dim, atom_dim, self.units), initializer=self.kernel_initializer, regularizer=self.kernel_regularizer, trainable=True, name="W", dtype=tf.float32, ) if self.use_bias: self.bias = self.add_weight( shape=(bond_dim, 1, self.units), initializer=self.bias_initializer, regularizer=self.bias_regularizer, trainable=True, name="b", dtype=tf.float32, ) self.built = True def call(self, inputs, training=False): adjacency, features = inputs # Aggregate information from neighbors x = tf.matmul(adjacency, features[:, None, :, :]) # Apply linear transformation x = tf.matmul(x, self.kernel) if self.use_bias: x += self.bias # Reduce bond types dim x_reduced = tf.reduce_sum(x, axis=1) # Apply non-linear transformation return self.activation(x_reduced) </code></pre></div> <hr /> <h2 id="build-the-encoder-and-decoder">Build the Encoder and Decoder</h2> The Encoder takes as input a molecule's graph adjacency matrix and feature matrix. These features are processed via a Graph Convolution layer, then are flattened and processed by several Dense layers to derive <code>z_mean</code> and <code>log_var</code>, the latent-space representation of the molecule. Graph Convolution layer: The relational graph convolution layer implements non-linearly transformed neighbourhood aggregations. We can define these layers as follows: <code>H_hat**(l+1) = σ(D_hat**(-1) * A_hat * H_hat**(l+1) * W**(l))</code> Where <code>σ</code> denotes the non-linear transformation (commonly a ReLU activation), <code>A</code> the adjacency tensor, <code>H_hat**(l)</code> the feature tensor at the <code>l-th</code> layer, <code>D_hat**(-1)</code> the inverse diagonal degree tensor of <code>A_hat</code>, and <code>W_hat**(l)</code> the trainable weight tensor at the <code>l-th</code> layer. Specifically, for each bond type (relation), the degree tensor expresses, in the diagonal, the number of bonds attached to each atom. Source: <a href="https://keras.io/examples/generative/wgan-graphs/">WGAN-GP with R-GCN for the generation of small molecular graphs</a>) The Decoder takes as input the latent-space representation and predicts the graph adjacency matrix and feature matrix of the corresponding molecules. <div class="codehilite"><pre><code>def get_encoder( gconv_units, latent_dim, adjacency_shape, feature_shape, dense_units, dropout_rate ): adjacency = keras.layers.Input(shape=adjacency_shape) features = keras.layers.Input(shape=feature_shape) # Propagate through one or more graph convolutional layers features_transformed = features for units in gconv_units: features_transformed = RelationalGraphConvLayer(units)( [adjacency, features_transformed] ) # Reduce 2-D representation of molecule to 1-D x = keras.layers.GlobalAveragePooling1D()(features_transformed) # Propagate through one or more densely connected layers for units in dense_units: x = layers.Dense(units, activation="relu")(x) x = layers.Dropout(dropout_rate)(x) z_mean = layers.Dense(latent_dim, dtype="float32", name="z_mean")(x) log_var = layers.Dense(latent_dim, dtype="float32", name="log_var")(x) encoder = keras.Model([adjacency, features], [z_mean, log_var], name="encoder") return encoder def get_decoder(dense_units, dropout_rate, latent_dim, adjacency_shape, feature_shape): latent_inputs = keras.Input(shape=(latent_dim,)) x = latent_inputs for units in dense_units: x = keras.layers.Dense(units, activation="tanh")(x) x = keras.layers.Dropout(dropout_rate)(x) # Map outputs of previous layer (x) to [continuous] adjacency tensors (x_adjacency) x_adjacency = keras.layers.Dense(tf.math.reduce_prod(adjacency_shape))(x) x_adjacency = keras.layers.Reshape(adjacency_shape)(x_adjacency) # Symmetrify tensors in the last two dimensions x_adjacency = (x_adjacency + tf.transpose(x_adjacency, (0, 1, 3, 2))) / 2 x_adjacency = keras.layers.Softmax(axis=1)(x_adjacency) # Map outputs of previous layer (x) to [continuous] feature tensors (x_features) x_features = keras.layers.Dense(tf.math.reduce_prod(feature_shape))(x) x_features = keras.layers.Reshape(feature_shape)(x_features) x_features = keras.layers.Softmax(axis=2)(x_features) decoder = keras.Model( latent_inputs, outputs=[x_adjacency, x_features], name="decoder" ) return decoder </code></pre></div> <hr /> <h2 id="build-the-sampling-layer">Build the Sampling layer</h2> <div class="codehilite"><pre><code>class Sampling(layers.Layer): def call(self, inputs): z_mean, z_log_var = inputs batch = tf.shape(z_log_var)[0] dim = tf.shape(z_log_var)[1] epsilon = tf.keras.backend.random_normal(shape=(batch, dim)) return z_mean + tf.exp(0.5 * z_log_var) * epsilon </code></pre></div> <hr /> <h2 id="build-the-vae">Build the VAE</h2> This model is trained to optimize four losses: <ul> <li>Categorical crossentropy</li> <li>KL divergence loss</li> <li>Property prediction loss</li> <li>Graph loss (gradient penalty)</li> </ul> The categorical crossentropy loss function measures the model's reconstruction accuracy. The Property prediction loss estimates the mean squared error between predicted and actual properties after running the latent representation through a property prediction model. The property prediction of the model is optimized via binary crossentropy. The gradient penalty is further guided by the model's property (QED) prediction. A gradient penalty is an alternative soft constraint on the 1-Lipschitz continuity as an improvement upon the gradient clipping scheme from the original neural network ("1-Lipschitz continuity" means that the norm of the gradient is at most 1 at every single point of the function). It adds a regularization term to the loss function. <div class="codehilite"><pre><code>class MoleculeGenerator(keras.Model): def __init__(self, encoder, decoder, max_len, **kwargs): super().__init__(**kwargs) self.encoder = encoder self.decoder = decoder self.property_prediction_layer = layers.Dense(1) self.max_len = max_len self.train_total_loss_tracker = keras.metrics.Mean(name="train_total_loss") self.val_total_loss_tracker = keras.metrics.Mean(name="val_total_loss") def train_step(self, data): adjacency_tensor, feature_tensor, qed_tensor = data[0] graph_real = [adjacency_tensor, feature_tensor] self.batch_size = tf.shape(qed_tensor)[0] with tf.GradientTape() as tape: z_mean, z_log_var, qed_pred, gen_adjacency, gen_features = self( graph_real, training=True ) graph_generated = [gen_adjacency, gen_features] total_loss = self._compute_loss( z_log_var, z_mean, qed_tensor, qed_pred, graph_real, graph_generated ) grads = tape.gradient(total_loss, self.trainable_weights) self.optimizer.apply_gradients(zip(grads, self.trainable_weights)) self.train_total_loss_tracker.update_state(total_loss) return {"loss": self.train_total_loss_tracker.result()} def _compute_loss( self, z_log_var, z_mean, qed_true, qed_pred, graph_real, graph_generated ): adjacency_real, features_real = graph_real adjacency_gen, features_gen = graph_generated adjacency_loss = tf.reduce_mean( tf.reduce_sum( keras.losses.categorical_crossentropy(adjacency_real, adjacency_gen), axis=(1, 2), ) ) features_loss = tf.reduce_mean( tf.reduce_sum( keras.losses.categorical_crossentropy(features_real, features_gen), axis=(1), ) ) kl_loss = -0.5 * tf.reduce_sum( 1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var), 1 ) kl_loss = tf.reduce_mean(kl_loss) property_loss = tf.reduce_mean( keras.losses.binary_crossentropy(qed_true, qed_pred) ) graph_loss = self._gradient_penalty(graph_real, graph_generated) return kl_loss + property_loss + graph_loss + adjacency_loss + features_loss def _gradient_penalty(self, graph_real, graph_generated): # Unpack graphs adjacency_real, features_real = graph_real adjacency_generated, features_generated = graph_generated # Generate interpolated graphs (adjacency_interp and features_interp) alpha = tf.random.uniform([self.batch_size]) alpha = tf.reshape(alpha, (self.batch_size, 1, 1, 1)) adjacency_interp = (adjacency_real * alpha) + (1 - alpha) * adjacency_generated alpha = tf.reshape(alpha, (self.batch_size, 1, 1)) features_interp = (features_real * alpha) + (1 - alpha) * features_generated # Compute the logits of interpolated graphs with tf.GradientTape() as tape: tape.watch(adjacency_interp) tape.watch(features_interp) _, _, logits, _, _ = self( [adjacency_interp, features_interp], training=True ) # Compute the gradients with respect to the interpolated graphs grads = tape.gradient(logits, [adjacency_interp, features_interp]) # Compute the gradient penalty grads_adjacency_penalty = (1 - tf.norm(grads[0], axis=1)) ** 2 grads_features_penalty = (1 - tf.norm(grads[1], axis=2)) ** 2 return tf.reduce_mean( tf.reduce_mean(grads_adjacency_penalty, axis=(-2, -1)) + tf.reduce_mean(grads_features_penalty, axis=(-1)) ) def inference(self, batch_size): z = tf.random.normal((batch_size, LATENT_DIM)) reconstruction_adjacency, reconstruction_features = model.decoder.predict(z) # obtain one-hot encoded adjacency tensor adjacency = tf.argmax(reconstruction_adjacency, axis=1) adjacency = tf.one_hot(adjacency, depth=BOND_DIM, axis=1) # Remove potential self-loops from adjacency adjacency = tf.linalg.set_diag(adjacency, tf.zeros(tf.shape(adjacency)[:-1])) # obtain one-hot encoded feature tensor features = tf.argmax(reconstruction_features, axis=2) features = tf.one_hot(features, depth=ATOM_DIM, axis=2) return [ graph_to_molecule([adjacency[i].numpy(), features[i].numpy()]) for i in range(batch_size) ] def call(self, inputs): z_mean, log_var = self.encoder(inputs) z = Sampling()([z_mean, log_var]) gen_adjacency, gen_features = self.decoder(z) property_pred = self.property_prediction_layer(z_mean) return z_mean, log_var, property_pred, gen_adjacency, gen_features </code></pre></div> <hr /> <h2 id="train-the-model">Train the model</h2> <div class="codehilite"><pre><code>vae_optimizer = tf.keras.optimizers.Adam(learning_rate=VAE_LR) encoder = get_encoder( gconv_units=[9], adjacency_shape=(BOND_DIM, NUM_ATOMS, NUM_ATOMS), feature_shape=(NUM_ATOMS, ATOM_DIM), latent_dim=LATENT_DIM, dense_units=[512], dropout_rate=0.0, ) decoder = get_decoder( dense_units=[128, 256, 512], dropout_rate=0.2, latent_dim=LATENT_DIM, adjacency_shape=(BOND_DIM, NUM_ATOMS, NUM_ATOMS), feature_shape=(NUM_ATOMS, ATOM_DIM), ) model = MoleculeGenerator(encoder, decoder, MAX_MOLSIZE) model.compile(vae_optimizer) history = model.fit([adjacency_tensor, feature_tensor, qed_tensor], epochs=EPOCHS) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Epoch 1/10 250/250 [==============================] - 24s 84ms/step - loss: 68958.3946 Epoch 2/10 250/250 [==============================] - 20s 79ms/step - loss: 68819.8421 Epoch 3/10 250/250 [==============================] - 20s 79ms/step - loss: 68830.6720 Epoch 4/10 250/250 [==============================] - 20s 79ms/step - loss: 68816.1486 Epoch 5/10 250/250 [==============================] - 20s 79ms/step - loss: 68825.9977 Epoch 6/10 250/250 [==============================] - 19s 78ms/step - loss: 68818.0771 Epoch 7/10 250/250 [==============================] - 19s 77ms/step - loss: 68815.8525 Epoch 8/10 250/250 [==============================] - 20s 78ms/step - loss: 68820.5459 Epoch 9/10 250/250 [==============================] - 21s 83ms/step - loss: 68806.9465 Epoch 10/10 250/250 [==============================] - 21s 84ms/step - loss: 68805.9879 </code></pre></div> </div> <hr /> <h2 id="inference">Inference</h2> We use our model to generate new valid molecules from different points of the latent space. <h3 id="generate-unique-molecules-with-the-model">Generate unique Molecules with the model</h3> <div class="codehilite"><pre><code>molecules = model.inference(1000) MolsToGridImage( [m for m in molecules if m is not None][:1000], molsPerRow=5, subImgSize=(260, 160) ) </code></pre></div> <img alt="png" src="/img/examples/generative/molecule_generation/molecule_generation_21_0.png" /> <h3 id="display-latent-space-clusters-with-respect-to-molecular-properties-qae">Display latent space clusters with respect to molecular properties (QAE)</h3> <hr /> <div class="codehilite"><pre><code>def plot_latent(vae, data, labels): # display a 2D plot of the property in the latent space z_mean, _ = vae.encoder.predict(data) plt.figure(figsize=(12, 10)) plt.scatter(z_mean[:, 0], z_mean[:, 1], c=labels) plt.colorbar() plt.xlabel("z[0]") plt.ylabel("z[1]") plt.show() plot_latent(model, [adjacency_tensor[:8000], feature_tensor[:8000]], qed_tensor[:8000]) </code></pre></div> <img alt="png" src="/img/examples/generative/molecule_generation/molecule_generation_23_0.png" /> <hr /> <h2 id="conclusion">Conclusion</h2> In this example, we combined model architectures from two papers, "Automatic chemical design using a data-driven continuous representation of molecules" from 2016 and the "MolGAN" paper from 2018. The former paper treats SMILES inputs as strings and seeks to generate molecule strings in SMILES format, while the later paper considers SMILES inputs as graphs (a combination of adjacency matrices and feature matrices) and seeks to generate molecules as graphs. This hybrid approach enables a new type of directed gradient-based search through chemical space. Example available on HuggingFace <table> <thead> <tr> <th style="text-align: center;">Trained Model</th> <th style="text-align: center;">Demo</th> </tr> </thead> <tbody> <tr> <td style="text-align: center;"><a href="https://huggingface.co/keras-io/drug-molecule-generation-with-VAE"><img alt="Generic badge" src="https://img.shields.io/badge/%F0%9F%A4%97%20Model-molecule%20generation%20with%20VAE-black.svg" /></a></td> <td style="text-align: center;"><a href="https://huggingface.co/spaces/keras-io/generating-drug-molecule-with-VAE"><img alt="Generic badge" src="https://img.shields.io/badge/%F0%9F%A4%97%20Spaces-molecule%20generation%20with%20VAE-black.svg" /></a></td> </tr> </tbody> </table> </div> <div class='k-outline'> <div class='k-outline-depth-1'> <a href='#drug-molecule-generation-with-vae'>Drug Molecule Generation with VAE</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#introduction'>Introduction</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#setup'>Setup</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#dataset'>Dataset</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#hyperparameters'>Hyperparameters</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#-generate-training-set'> Generate training set</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#build-the-encoder-and-decoder'>Build the Encoder and Decoder</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#build-the-sampling-layer'>Build the Sampling layer</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#build-the-vae'>Build the VAE</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#train-the-model'>Train the model</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#inference'>Inference</a> </div> <div class='k-outline-depth-3'> <a href='#generate-unique-molecules-with-the-model'>Generate unique Molecules with the model</a> </div> <div class='k-outline-depth-3'> <a href='#display-latent-space-clusters-with-respect-to-molecular-properties-qae'>Display latent space clusters with respect to molecular properties (QAE)</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#conclusion'>Conclusion</a> </div> </div> </div> </div> </div> </body> <footer style="float: left; width: 100%; padding: 1em; border-top: solid 1px #bbb;"> <a href="https://policies.google.com/terms">Terms</a> | <a href="https://policies.google.com/privacy">Privacy</a> </footer> </html>