train.py

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)