--- a
+++ b/train.py
@@ -0,0 +1,462 @@
+import os
+import time
+import random
+import pickle
+import argparse
+import os.path as osp
+
+import torch
+import torch.utils.data
+from torch import nn
+from torch_geometric.loader import DataLoader
+
+import wandb
+from rdkit import RDLogger
+
+torch.set_num_threads(5)
+RDLogger.DisableLog('rdApp.*')
+
+from src.util.utils import *
+from src.model.models import Generator, Discriminator, simple_disc
+from src.data.dataset import DruggenDataset
+from src.data.utils import get_encoders_decoders, load_molecules
+from src.model.loss import discriminator_loss, generator_loss
+
+class Train(object):
+ """Trainer for DrugGEN."""
+
+ def __init__(self, config):
+ if config.set_seed:
+ np.random.seed(config.seed)
+ random.seed(config.seed)
+ torch.manual_seed(config.seed)
+ torch.cuda.manual_seed_all(config.seed)
+
+ torch.backends.cudnn.deterministic = True
+ torch.backends.cudnn.benchmark = False
+
+ os.environ["PYTHONHASHSEED"] = str(config.seed)
+
+ print(f'Using seed {config.seed}')
+
+ self.device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')
+
+ # Initialize configurations
+ self.submodel = config.submodel
+
+ # Data loader.
+ self.raw_file = config.raw_file # SMILES containing text file for dataset.
+ # Write the full path to file.
+ self.drug_raw_file = config.drug_raw_file # SMILES containing text file for second dataset.
+ # Write the full path to file.
+
+ # Automatically infer dataset file names from raw file names
+ raw_file_basename = osp.basename(self.raw_file)
+ drug_raw_file_basename = osp.basename(self.drug_raw_file)
+
+ # Get the base name without extension and add max_atom to it
+ self.max_atom = config.max_atom # Model is based on one-shot generation.
+ raw_file_base = os.path.splitext(raw_file_basename)[0]
+ drug_raw_file_base = os.path.splitext(drug_raw_file_basename)[0]
+
+ # Change extension from .smi to .pt and add max_atom to the filename
+ self.dataset_file = f"{raw_file_base}{self.max_atom}.pt"
+ self.drugs_dataset_file = f"{drug_raw_file_base}{self.max_atom}.pt"
+
+ self.mol_data_dir = config.mol_data_dir # Directory where the dataset files are stored.
+ self.drug_data_dir = config.drug_data_dir # Directory where the drug dataset files are stored.
+ self.dataset_name = self.dataset_file.split(".")[0]
+ self.drugs_dataset_name = self.drugs_dataset_file.split(".")[0]
+ self.features = config.features # Small model uses atom types as node features. (Boolean, False uses atom types only.)
+ # Additional node features can be added. Please check new_dataloarder.py Line 102.
+ self.batch_size = config.batch_size # Batch size for training.
+
+ self.parallel = config.parallel
+
+ # Get atom and bond encoders/decoders
+ atom_encoder, atom_decoder, bond_encoder, bond_decoder = get_encoders_decoders(
+ self.raw_file,
+ self.drug_raw_file,
+ self.max_atom
+ )
+ self.atom_encoder = atom_encoder
+ self.atom_decoder = atom_decoder
+ self.bond_encoder = bond_encoder
+ self.bond_decoder = bond_decoder
+
+ self.dataset = DruggenDataset(self.mol_data_dir,
+ self.dataset_file,
+ self.raw_file,
+ self.max_atom,
+ self.features,
+ atom_encoder=atom_encoder,
+ atom_decoder=atom_decoder,
+ bond_encoder=bond_encoder,
+ bond_decoder=bond_decoder)
+
+ self.loader = DataLoader(self.dataset,
+ shuffle=True,
+ batch_size=self.batch_size,
+ drop_last=True) # PyG dataloader for the GAN.
+
+ self.drugs = DruggenDataset(self.drug_data_dir,
+ self.drugs_dataset_file,
+ self.drug_raw_file,
+ self.max_atom,
+ self.features,
+ atom_encoder=atom_encoder,
+ atom_decoder=atom_decoder,
+ bond_encoder=bond_encoder,
+ bond_decoder=bond_decoder)
+
+ self.drugs_loader = DataLoader(self.drugs,
+ shuffle=True,
+ batch_size=self.batch_size,
+ drop_last=True) # PyG dataloader for the second GAN.
+
+ self.m_dim = len(self.atom_decoder) if not self.features else int(self.loader.dataset[0].x.shape[1]) # Atom type dimension.
+ self.b_dim = len(self.bond_decoder) # Bond type dimension.
+ self.vertexes = int(self.loader.dataset[0].x.shape[0]) # Number of nodes in the graph.
+
+ # Model configurations.
+ self.act = config.act
+ self.lambda_gp = config.lambda_gp
+ self.dim = config.dim
+ self.depth = config.depth
+ self.heads = config.heads
+ self.mlp_ratio = config.mlp_ratio
+ self.ddepth = config.ddepth
+ self.ddropout = config.ddropout
+
+ # Training configurations.
+ self.epoch = config.epoch
+ self.g_lr = config.g_lr
+ self.d_lr = config.d_lr
+ self.dropout = config.dropout
+ self.beta1 = config.beta1
+ self.beta2 = config.beta2
+
+ # Directories.
+ self.log_dir = config.log_dir
+ self.sample_dir = config.sample_dir
+ self.model_save_dir = config.model_save_dir
+
+ # Step size.
+ self.log_step = config.log_sample_step
+
+ # resume training
+ self.resume = config.resume
+ self.resume_epoch = config.resume_epoch
+ self.resume_iter = config.resume_iter
+ self.resume_directory = config.resume_directory
+
+ # wandb configuration
+ self.use_wandb = config.use_wandb
+ self.online = config.online
+ self.exp_name = config.exp_name
+
+ # Arguments for the model.
+ self.arguments = "{}_{}_glr{}_dlr{}_dim{}_depth{}_heads{}_batch{}_epoch{}_dataset{}_dropout{}".format(self.exp_name, self.submodel, self.g_lr, self.d_lr, self.dim, self.depth, self.heads, self.batch_size, self.epoch, self.dataset_name, self.dropout)
+
+ self.build_model(self.model_save_dir, self.arguments)
+
+
+ def build_model(self, model_save_dir, arguments):
+ """Create generators and discriminators."""
+
+ ''' Generator is based on Transformer Encoder:
+
+ @ g_conv_dim: Dimensions for MLP layers before Transformer Encoder
+ @ vertexes: maximum length of generated molecules (atom length)
+ @ b_dim: number of bond types
+ @ m_dim: number of atom types (or number of features used)
+ @ dropout: dropout possibility
+ @ dim: Hidden dimension of Transformer Encoder
+ @ depth: Transformer layer number
+ @ heads: Number of multihead-attention heads
+ @ mlp_ratio: Read-out layer dimension of Transformer
+ @ drop_rate: depricated
+ @ tra_conv: Whether module creates output for TransformerConv discriminator
+ '''
+ self.G = Generator(self.act,
+ self.vertexes,
+ self.b_dim,
+ self.m_dim,
+ self.dropout,
+ dim=self.dim,
+ depth=self.depth,
+ heads=self.heads,
+ mlp_ratio=self.mlp_ratio)
+
+ ''' Discriminator implementation with Transformer Encoder:
+
+ @ act: Activation function for MLP
+ @ vertexes: maximum length of generated molecules (molecule length)
+ @ b_dim: number of bond types
+ @ m_dim: number of atom types (or number of features used)
+ @ dropout: dropout possibility
+ @ dim: Hidden dimension of Transformer Encoder
+ @ depth: Transformer layer number
+ @ heads: Number of multihead-attention heads
+ @ mlp_ratio: Read-out layer dimension of Transformer'''
+
+ self.D = Discriminator(self.act,
+ self.vertexes,
+ self.b_dim,
+ self.m_dim,
+ self.ddropout,
+ dim=self.dim,
+ depth=self.ddepth,
+ heads=self.heads,
+ mlp_ratio=self.mlp_ratio)
+
+ self.g_optimizer = torch.optim.AdamW(self.G.parameters(), self.g_lr, [self.beta1, self.beta2])
+ self.d_optimizer = torch.optim.AdamW(self.D.parameters(), self.d_lr, [self.beta1, self.beta2])
+
+ network_path = os.path.join(model_save_dir, arguments)
+ self.print_network(self.G, 'G', network_path)
+ self.print_network(self.D, 'D', network_path)
+
+ if self.parallel and torch.cuda.device_count() > 1:
+ print(f"Using {torch.cuda.device_count()} GPUs!")
+ self.G = nn.DataParallel(self.G)
+ self.D = nn.DataParallel(self.D)
+
+ self.G.to(self.device)
+ self.D.to(self.device)
+
+ def print_network(self, model, name, save_dir):
+ """Print out the network information."""
+ num_params = 0
+ for p in model.parameters():
+ num_params += p.numel()
+
+ if not os.path.exists(save_dir):
+ os.makedirs(save_dir)
+
+ network_path = os.path.join(save_dir, "{}_modules.txt".format(name))
+ with open(network_path, "w+") as file:
+ for module in model.modules():
+ file.write(f"{module.__class__.__name__}:\n")
+ print(module.__class__.__name__)
+ for n, param in module.named_parameters():
+ if param is not None:
+ file.write(f" - {n}: {param.size()}\n")
+ print(f" - {n}: {param.size()}")
+ break
+ file.write(f"Total number of parameters: {num_params}\n")
+ print(f"Total number of parameters: {num_params}\n\n")
+
+ def restore_model(self, epoch, iteration, model_directory):
+ """Restore the trained generator and discriminator."""
+ print('Loading the trained models from epoch / iteration {}-{}...'.format(epoch, iteration))
+
+ G_path = os.path.join(model_directory, '{}-{}-G.ckpt'.format(epoch, iteration))
+ D_path = os.path.join(model_directory, '{}-{}-D.ckpt'.format(epoch, iteration))
+ self.G.load_state_dict(torch.load(G_path, map_location=lambda storage, loc: storage))
+ self.D.load_state_dict(torch.load(D_path, map_location=lambda storage, loc: storage))
+
+ def save_model(self, model_directory, idx,i):
+ G_path = os.path.join(model_directory, '{}-{}-G.ckpt'.format(idx+1,i+1))
+ D_path = os.path.join(model_directory, '{}-{}-D.ckpt'.format(idx+1,i+1))
+ torch.save(self.G.state_dict(), G_path)
+ torch.save(self.D.state_dict(), D_path)
+
+ def reset_grad(self):
+ """Reset the gradient buffers."""
+ self.g_optimizer.zero_grad()
+ self.d_optimizer.zero_grad()
+
+ def train(self, config):
+ ''' Training Script starts from here'''
+ if self.use_wandb:
+ mode = 'online' if self.online else 'offline'
+ else:
+ mode = 'disabled'
+ kwargs = {'name': self.exp_name, 'project': 'druggen', 'config': config,
+ 'settings': wandb.Settings(_disable_stats=True), 'reinit': True, 'mode': mode, 'save_code': True}
+ wandb.init(**kwargs)
+
+ wandb.save(os.path.join(self.model_save_dir, self.arguments, "G_modules.txt"))
+ wandb.save(os.path.join(self.model_save_dir, self.arguments, "D_modules.txt"))
+
+ self.model_directory = os.path.join(self.model_save_dir, self.arguments)
+ self.sample_directory = os.path.join(self.sample_dir, self.arguments)
+ self.log_path = os.path.join(self.log_dir, "{}.txt".format(self.arguments))
+ if not os.path.exists(self.model_directory):
+ os.makedirs(self.model_directory)
+ if not os.path.exists(self.sample_directory):
+ os.makedirs(self.sample_directory)
+
+ # smiles data for metrics calculation.
+ drug_smiles = [line for line in open(self.drug_raw_file, 'r').read().splitlines()]
+ drug_mols = [Chem.MolFromSmiles(smi) for smi in drug_smiles]
+ drug_vecs = [AllChem.GetMorganFingerprintAsBitVect(x, 2, nBits=1024) for x in drug_mols if x is not None]
+
+ if self.resume:
+ self.restore_model(self.resume_epoch, self.resume_iter, self.resume_directory)
+
+ # Start training.
+ print('Start training...')
+ self.start_time = time.time()
+ for idx in range(self.epoch):
+ # =================================================================================== #
+ # 1. Preprocess input data #
+ # =================================================================================== #
+ # Load the data
+ dataloader_iterator = iter(self.drugs_loader)
+
+ wandb.log({"epoch": idx})
+
+ for i, data in enumerate(self.loader):
+ try:
+ drugs = next(dataloader_iterator)
+ except StopIteration:
+ dataloader_iterator = iter(self.drugs_loader)
+ drugs = next(dataloader_iterator)
+
+ wandb.log({"iter": i})
+
+ # Preprocess both dataset
+ real_graphs, a_tensor, x_tensor = load_molecules(
+ data=data,
+ batch_size=self.batch_size,
+ device=self.device,
+ b_dim=self.b_dim,
+ m_dim=self.m_dim,
+ )
+
+ drug_graphs, drugs_a_tensor, drugs_x_tensor = load_molecules(
+ data=drugs,
+ batch_size=self.batch_size,
+ device=self.device,
+ b_dim=self.b_dim,
+ m_dim=self.m_dim,
+ )
+
+ # Training configuration.
+ GEN_node = x_tensor # Generator input node features (annotation matrix of real molecules)
+ GEN_edge = a_tensor # Generator input edge features (adjacency matrix of real molecules)
+ if self.submodel == "DrugGEN":
+ DISC_node = drugs_x_tensor # Discriminator input node features (annotation matrix of drug molecules)
+ DISC_edge = drugs_a_tensor # Discriminator input edge features (adjacency matrix of drug molecules)
+ elif self.submodel == "NoTarget":
+ DISC_node = x_tensor # Discriminator input node features (annotation matrix of real molecules)
+ DISC_edge = a_tensor # Discriminator input edge features (adjacency matrix of real molecules)
+
+ # =================================================================================== #
+ # 2. Train the GAN #
+ # =================================================================================== #
+
+ loss = {}
+ self.reset_grad()
+ # Compute discriminator loss.
+ node, edge, d_loss = discriminator_loss(self.G,
+ self.D,
+ DISC_edge,
+ DISC_node,
+ GEN_edge,
+ GEN_node,
+ self.batch_size,
+ self.device,
+ self.lambda_gp)
+ d_total = d_loss
+ wandb.log({"d_loss": d_total.item()})
+
+ loss["d_total"] = d_total.item()
+ d_total.backward()
+ self.d_optimizer.step()
+
+ self.reset_grad()
+
+ # Compute generator loss.
+ generator_output = generator_loss(self.G,
+ self.D,
+ GEN_edge,
+ GEN_node,
+ self.batch_size)
+ g_loss, node, edge, node_sample, edge_sample = generator_output
+ g_total = g_loss
+ wandb.log({"g_loss": g_total.item()})
+
+ loss["g_total"] = g_total.item()
+ g_total.backward()
+ self.g_optimizer.step()
+
+ # Logging.
+ if (i+1) % self.log_step == 0:
+ logging(self.log_path, self.start_time, i, idx, loss, self.sample_directory,
+ drug_smiles,edge_sample, node_sample, self.dataset.matrices2mol,
+ self.dataset_name, a_tensor, x_tensor, drug_vecs)
+
+ mol_sample(self.sample_directory, edge_sample.detach(), node_sample.detach(),
+ idx, i, self.dataset.matrices2mol, self.dataset_name)
+ print("samples saved at epoch {} and iteration {}".format(idx,i))
+
+ self.save_model(self.model_directory, idx, i)
+ print("model saved at epoch {} and iteration {}".format(idx,i))
+
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser()
+
+ # Data configuration.
+ parser.add_argument('--raw_file', type=str, required=True)
+ parser.add_argument('--drug_raw_file', type=str, required=False, help='Required for DrugGEN model, optional for NoTarget')
+ parser.add_argument('--drug_data_dir', type=str, default='data')
+ parser.add_argument('--mol_data_dir', type=str, default='data')
+ parser.add_argument('--features', action='store_true', help='features dimension for nodes')
+
+ # Model configuration.
+ parser.add_argument('--submodel', type=str, default="DrugGEN", help="Chose model subtype: DrugGEN, NoTarget", choices=['DrugGEN', 'NoTarget'])
+ parser.add_argument('--act', type=str, default="relu", help="Activation function for the model.", choices=['relu', 'tanh', 'leaky', 'sigmoid'])
+ parser.add_argument('--max_atom', type=int, default=45, help='Max atom number for molecules must be specified.')
+ parser.add_argument('--dim', type=int, default=128, help='Dimension of the Transformer Encoder model for the GAN.')
+ parser.add_argument('--depth', type=int, default=1, help='Depth of the Transformer model from the GAN.')
+ parser.add_argument('--ddepth', type=int, default=1, help='Depth of the Transformer model from the discriminator.')
+ parser.add_argument('--heads', type=int, default=8, help='Number of heads for the MultiHeadAttention module from the GAN.')
+ parser.add_argument('--mlp_ratio', type=int, default=3, help='MLP ratio for the Transformer.')
+ parser.add_argument('--dropout', type=float, default=0., help='dropout rate')
+ parser.add_argument('--ddropout', type=float, default=0., help='dropout rate for the discriminator')
+ parser.add_argument('--lambda_gp', type=float, default=10, help='Gradient penalty lambda multiplier for the GAN.')
+
+ # Training configuration.
+ parser.add_argument('--batch_size', type=int, default=128, help='Batch size for the training.')
+ parser.add_argument('--epoch', type=int, default=10, help='Epoch number for Training.')
+ parser.add_argument('--g_lr', type=float, default=0.00001, help='learning rate for G')
+ parser.add_argument('--d_lr', type=float, default=0.00001, help='learning rate for D')
+ parser.add_argument('--beta1', type=float, default=0.9, help='beta1 for Adam optimizer')
+ parser.add_argument('--beta2', type=float, default=0.999, help='beta2 for Adam optimizer')
+ parser.add_argument('--log_dir', type=str, default='experiments/logs')
+ parser.add_argument('--sample_dir', type=str, default='experiments/samples')
+ parser.add_argument('--model_save_dir', type=str, default='experiments/models')
+ parser.add_argument('--log_sample_step', type=int, default=1000, help='step size for sampling during training')
+
+ # Resume training.
+ parser.add_argument('--resume', type=bool, default=False, help='resume training')
+ parser.add_argument('--resume_epoch', type=int, default=None, help='resume training from this epoch')
+ parser.add_argument('--resume_iter', type=int, default=None, help='resume training from this step')
+ parser.add_argument('--resume_directory', type=str, default=None, help='load pretrained weights from this directory')
+
+ # Seed configuration.
+ parser.add_argument('--set_seed', action='store_true', help='set seed for reproducibility')
+ parser.add_argument('--seed', type=int, default=1, help='seed for reproducibility')
+
+ # wandb configuration.
+ parser.add_argument('--use_wandb', action='store_true', help='use wandb for logging')
+ parser.add_argument('--online', action='store_true', help='use wandb online')
+ parser.add_argument('--exp_name', type=str, default='druggen', help='experiment name')
+ parser.add_argument('--parallel', action='store_true', help='Parallelize training')
+
+ config = parser.parse_args()
+
+ # Check if drug_raw_file is provided when using DrugGEN model
+ if config.submodel == "DrugGEN" and not config.drug_raw_file:
+ parser.error("--drug_raw_file is required when using DrugGEN model")
+
+ # If using NoTarget model and drug_raw_file is not provided, use a dummy file
+ if config.submodel == "NoTarget" and not config.drug_raw_file:
+ config.drug_raw_file = "data/akt_train.smi" # Use a reference file for NoTarget model (AKT) (not used for training for ease of use and encoder/decoder's)
+
+ trainer = Train(config)
+ trainer.train(config)
Datasets
Models
