dscript.commands.train — D-SCRIPT v1.0-beta documentation

# Source code for dscript.commands.train

    """
    Train a new model.
    """

    import sys
    import argparse
    import h5py
    import datetime
    import subprocess as sp
    import numpy as np
    import pandas as pd
    import gzip as gz
    from tqdm import tqdm

    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    import torch.optim as optim
    from torch.autograd import Variable
    from torch.utils.data import IterableDataset, DataLoader
    from sklearn.metrics import average_precision_score as average_precision

    import dscript
    from dscript.utils import PairedDataset, collate_paired_sequences
    from dscript.models.embedding import (
        IdentityEmbed,
        FullyConnectedEmbed,
    )
    from dscript.models.contact import ContactCNN
    from dscript.models.interaction import ModelInteraction

    def add_args(parser):
        """
        Create parser for command line utility.

        :meta private:
        """

        data_grp = parser.add_argument_group("Data")
        proj_grp = parser.add_argument_group("Projection Module")
        contact_grp = parser.add_argument_group("Contact Module")
        inter_grp = parser.add_argument_group("Interaction Module")
        train_grp = parser.add_argument_group("Training")
        misc_grp = parser.add_argument_group("Output and Device")

        # Data
        data_grp.add_argument("--train", help="Training data", required=True)
        data_grp.add_argument("--val", help="Validation data", required=True)
        data_grp.add_argument("--embedding", help="h5 file with embedded sequences", required=True)
        data_grp.add_argument(
            "--augment",
            action="store_true",
            help="Set flag to augment data by adding (B A) for all pairs (A B)",
        )

        # Embedding model
        proj_grp.add_argument(
            "--projection-dim",
            type=int,
            default=100,
            help="Dimension of embedding projection layer (default: 100)",
        )
        proj_grp.add_argument(
            "--dropout-p",
            type=float,
            default=0.5,
            help="Parameter p for embedding dropout layer (default: 0.5)",
        )

        # Contact model
        contact_grp.add_argument(
            "--hidden-dim",
            type=int,
            default=50,
            help="Number of hidden units for comparison layer in contact prediction (default: 50)",
        )
        contact_grp.add_argument(
            "--kernel-width",
            type=int,
            default=7,
            help="Width of convolutional filter for contact prediction (default: 7)",
        )

        # Interaction Model
        inter_grp.add_argument(
            "--use-w",
            action="store_true",
            help="Use weight matrix in interaction prediction model",
        )
        inter_grp.add_argument(
            "--pool-width",
            type=int,
            default=9,
            help="Size of max-pool in interaction model (default: 9)",
        )

        # Training
        train_grp.add_argument(
            "--negative-ratio",
            type=int,
            default=10,
            help="Number of negative training samples for each positive training sample (default: 10)",
        )
        train_grp.add_argument(
            "--epoch-scale",
            type=int,
            default=5,
            help="Report heldout performance every this many epochs (default: 5)",
        )
        train_grp.add_argument("--num-epochs", type=int, default=100, help="Number of epochs (default: 100)")
        train_grp.add_argument("--batch-size", type=int, default=25, help="Minibatch size (default: 25)")
        train_grp.add_argument("--weight-decay", type=float, default=0, help="L2 regularization (default: 0)")
        train_grp.add_argument("--lr", type=float, default=0.001, help="Learning rate (default: 0.001)")
        train_grp.add_argument(
            "--lambda",
            dest="lambda_",
            type=float,
            default=0.35,
            help="Weight on the similarity objective (default: 0.35)",
        )

        # Output
        misc_grp.add_argument("-o", "--outfile", help="Output file path (default: stdout)")
        misc_grp.add_argument("--save-prefix", help="Path prefix for saving models")
        misc_grp.add_argument("-d", "--device", type=int, default=-1, help="Compute device to use")
        misc_grp.add_argument("--checkpoint", help="Checkpoint model to start training from")

        return parser

    def predict_interaction(model, n0, n1, tensors, use_cuda):
        """
        Predict whether a list of protein pairs will interact.

        :param model: Model to be trained
        :type model: dscript.models.interaction.ModelInteraction
        :param n0: First protein names
        :type n0: list[str]
        :param n1: Second protein names
        :type n1: list[str]
        :param tensors: Dictionary of protein names to embeddings
        :type tensors: dict[str, torch.Tensor]
        :param use_cuda: Whether to use GPU
        :type use_cuda: bool
        """

        b = len(n0)

        p_hat = []
        for i in range(b):
            z_a = tensors[n0[i]]
            z_b = tensors[n1[i]]
            if use_cuda:
                z_a = z_a.cuda()
                z_b = z_b.cuda()

            p_hat.append(model.predict(z_a, z_b))
        p_hat = torch.stack(p_hat, 0)
        return p_hat

    def predict_cmap_interaction(model, n0, n1, tensors, use_cuda):
        """
        Predict whether a list of protein pairs will interact, as well as their contact map.

        :param model: Model to be trained
        :type model: dscript.models.interaction.ModelInteraction
        :param n0: First protein names
        :type n0: list[str]
        :param n1: Second protein names
        :type n1: list[str]
        :param tensors: Dictionary of protein names to embeddings
        :type tensors: dict[str, torch.Tensor]
        :param use_cuda: Whether to use GPU
        :type use_cuda: bool
        """

        b = len(n0)

        p_hat = []
        c_map_mag = []
        for i in range(b):
            z_a = tensors[n0[i]]
            z_b = tensors[n1[i]]
            if use_cuda:
                z_a = z_a.cuda()
                z_b = z_b.cuda()

            cm, ph = model.map_predict(z_a, z_b)
            p_hat.append(ph)
            c_map_mag.append(torch.mean(cm))
        p_hat = torch.stack(p_hat, 0)
        c_map_mag = torch.stack(c_map_mag, 0)
        return c_map_mag, p_hat

    def interaction_grad(model, n0, n1, y, tensors, use_cuda, weight=0.35):
        """
        Compute gradient and backpropagate loss for a batch.

        :param model: Model to be trained
        :type model: dscript.models.interaction.ModelInteraction
        :param n0: First protein names
        :type n0: list[str]
        :param n1: Second protein names
        :type n1: list[str]
        :param y: Interaction labels
        :type y: torch.Tensor
        :param tensors: Dictionary of protein names to embeddings
        :type tensors: dict[str, torch.Tensor]
        :param use_cuda: Whether to use GPU
        :type use_cuda: bool
        :param weight: Weight on the contact map magnitude objective. BCE loss is :math:`1 - \\text{weight}`.
        :type weight: float

        :return: (Loss, number correct, mean square error, batch size)
        :rtype: (torch.Tensor, int, torch.Tensor, int)
        """

        c_map_mag, p_hat = predict_cmap_interaction(model, n0, n1, tensors, use_cuda)
        if use_cuda:
            y = y.cuda()
        y = Variable(y)

        bce_loss = F.binary_cross_entropy(p_hat.float(), y.float())
        cmap_loss = torch.mean(c_map_mag)
        loss = (weight * bce_loss) + ((1 - weight) * cmap_loss)
        b = len(p_hat)

        # backprop loss
        loss.backward()

        if use_cuda:
            y = y.cpu()
            p_hat = p_hat.cpu()

        with torch.no_grad():
            guess_cutoff = 0.5
            p_hat = p_hat.float()
            p_guess = (guess_cutoff * torch.ones(b) < p_hat).float()
            y = y.float()
            correct = torch.sum(p_guess == y).item()
            mse = torch.mean((y.float() - p_hat) ** 2).item()

        return loss, correct, mse, b

    def interaction_eval(model, test_iterator, tensors, use_cuda):
        """
        Evaluate test data set performance.

        :param model: Model to be trained
        :type model: dscript.models.interaction.ModelInteraction
        :param test_iterator: Test data iterator
        :type test_iterator: torch.utils.data.DataLoader
        :param tensors: Dictionary of protein names to embeddings
        :type tensors: dict[str, torch.Tensor]
        :param use_cuda: Whether to use GPU
        :type use_cuda: bool

        :return: (Loss, number correct, mean square error, precision, recall, F1 Score, AUPR)
        :rtype: (torch.Tensor, int, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor)
        """
        p_hat = []
        true_y = []

        for n0, n1, y in test_iterator:
            p_hat.append(predict_interaction(model, n0, n1, tensors, use_cuda))
            true_y.append(y)

        y = torch.cat(true_y, 0)
        p_hat = torch.cat(p_hat, 0)

        if use_cuda:
            y.cuda()
            p_hat = torch.Tensor([x.cuda() for x in p_hat])
            p_hat.cuda()

        loss = F.binary_cross_entropy(p_hat.float(), y.float()).item()
        b = len(y)

        with torch.no_grad():
            guess_cutoff = torch.Tensor([0.5]).float()
            p_hat = p_hat.float()
            y = y.float()
            p_guess = (guess_cutoff * torch.ones(b) < p_hat).float()
            correct = torch.sum(p_guess == y).item()
            mse = torch.mean((y.float() - p_hat) ** 2).item()

            tp = torch.sum(y * p_hat).item()
            pr = tp / torch.sum(p_hat).item()
            re = tp / torch.sum(y).item()
            f1 = 2 * pr * re / (pr + re)

        y = y.cpu().numpy()
        p_hat = p_hat.data.cpu().numpy()

        aupr = average_precision(y, p_hat)

        return loss, correct, mse, pr, re, f1, aupr

    def main(args):
        """
        Run training from arguments.

        :meta private:
        """

        output = args.outfile
        if output is None:
            output = sys.stdout
        else:
            output = open(output, "w")

        print(f'# Called as: {" ".join(sys.argv)}', file=output)
        if output is not sys.stdout:
            print(f'Called as: {" ".join(sys.argv)}')

        # Set device
        device = args.device
        use_cuda = (device >= 0) and torch.cuda.is_available()
        if use_cuda:
            torch.cuda.set_device(device)
            print(
                f"# Using CUDA device {device} - {torch.cuda.get_device_name(device)}",
                file=output,
            )
        else:
            print("# Using CPU", file=output)
            device = "cpu"

        batch_size = args.batch_size

        train_fi = args.train
        test_fi = args.val
        augment = args.augment
        embedding_h5 = args.embedding
        h5fi = h5py.File(embedding_h5, "r")

        print(f"# Loading training pairs from {train_fi}...", file=output)
        output.flush()

        train_df = pd.read_csv(train_fi, sep="\t", header=None)
        if augment:
            train_n0 = pd.concat((train_df[0], train_df[1]), axis=0)
            train_n1 = pd.concat((train_df[1], train_df[0]), axis=0)
            train_y = torch.from_numpy(pd.concat((train_df[2], 1 - train_df[2])).values)
        else:
            train_n0, train_n1 = train_df[0], train_df[1]
            train_y = torch.from_numpy(train_df[2].values)

        print(f"# Loading testing pairs from {test_fi}...", file=output)
        output.flush()

        test_df = pd.read_csv(test_fi, sep="\t", header=None)
        if augment:
            test_n0 = pd.concat((test_df[0], test_df[1]), axis=0)
            test_n1 = pd.concat((test_df[1], test_df[0]), axis=0)
            test_y = torch.from_numpy(pd.concat((test_df[2], 1 - test_df[2])).values)
        else:
            test_n0, test_n1 = test_df[0], test_df[1]
            test_y = torch.from_numpy(test_df[2].values)
        output.flush()

        train_pairs = PairedDataset(train_n0, train_n1, train_y)
        pairs_train_iterator = torch.utils.data.DataLoader(
            train_pairs,
            batch_size=batch_size,
            collate_fn=collate_paired_sequences,
            shuffle=True,
        )

        test_pairs = PairedDataset(test_n0, test_n1, test_y)
        pairs_test_iterator = torch.utils.data.DataLoader(
            test_pairs,
            batch_size=batch_size,
            collate_fn=collate_paired_sequences,
            shuffle=True,
        )

        output.flush()

        print(f"# Loading embeddings", file=output)
        tensors = {}
        all_proteins = set(train_n0).union(set(train_n1)).union(set(test_n0)).union(set(test_n1))
        for prot_name in tqdm(all_proteins):
            tensors[prot_name] = torch.from_numpy(h5fi[prot_name][:, :])

        use_cuda = (args.device > -1) and torch.cuda.is_available()

        if args.checkpoint is None:

            projection_dim = args.projection_dim
            dropout_p = args.dropout_p
            embedding = FullyConnectedEmbed(6165, projection_dim, dropout=dropout_p)
            print("# Initializing embedding model with:", file=output)
            print(f"\tprojection_dim: {projection_dim}", file=output)
            print(f"\tdropout_p: {dropout_p}", file=output)

            # Create contact model
            hidden_dim = args.hidden_dim
            kernel_width = args.kernel_width
            print("# Initializing contact model with:", file=output)
            print(f"\thidden_dim: {hidden_dim}", file=output)
            print(f"\tkernel_width: {kernel_width}", file=output)

            contact = ContactCNN(projection_dim, hidden_dim, kernel_width)

            # Create the full model
            use_W = args.use_w
            pool_width = args.pool_width
            print("# Initializing interaction model with:", file=output)
            print(f"\tpool_width: {pool_width}", file=output)
            print(f"\tuse_w: {use_W}", file=output)
            model = ModelInteraction(embedding, contact, use_W=use_W, pool_size=pool_width)

            print(model, file=output)

        else:
            print("# Loading model from checkpoint {}".format(args.checkpoint), file=output)
            model = torch.load(args.checkpoint)
            model.use_cuda = use_cuda

        if use_cuda:
            model = model.cuda()

        # Train the model
        lr = args.lr
        wd = args.weight_decay
        num_epochs = args.num_epochs
        batch_size = args.batch_size
        report_steps = args.epoch_scale
        inter_weight = args.lambda_
        cmap_weight = 1 - inter_weight
        digits = int(np.floor(np.log10(num_epochs))) + 1
        save_prefix = args.save_prefix
        if save_prefix is None:
            save_prefix = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M")

        params = [p for p in model.parameters() if p.requires_grad]
        optim = torch.optim.Adam(params, lr=lr, weight_decay=wd)

        print(f'# Using save prefix "{save_prefix}"', file=output)
        print(f"# Training with Adam: lr={lr}, weight_decay={wd}", file=output)
        print(f"\tnum_epochs: {num_epochs}", file=output)
        print(f"\tepoch_scale: {report_steps}", file=output)
        print(f"\tbatch_size: {batch_size}", file=output)
        print(f"\tinteraction weight: {inter_weight}", file=output)
        print(f"\tcontact map weight: {cmap_weight}", file=output)
        output.flush()

        batch_report_fmt = "# [{}/{}] training {:.1%}: Loss={:.6}, Accuracy={:.3%}, MSE={:.6}"
        epoch_report_fmt = "# Finished Epoch {}/{}: Loss={:.6}, Accuracy={:.3%}, MSE={:.6}, Precision={:.6}, Recall={:.6}, F1={:.6}, AUPR={:.6}"

        N = len(pairs_train_iterator) * batch_size
        for epoch in range(num_epochs):

            model.train()

            n = 0
            loss_accum = 0
            acc_accum = 0
            mse_accum = 0

            # Train batches
            for (z0, z1, y) in tqdm(pairs_train_iterator, desc=f"Epoch {epoch+1}/{num_epochs}",total=len(pairs_train_iterator)):

                loss, correct, mse, b = interaction_grad(model, z0, z1, y, tensors, use_cuda, weight=inter_weight)

                n += b
                delta = b * (loss - loss_accum)
                loss_accum += delta / n

                delta = correct - b * acc_accum
                acc_accum += delta / n

                delta = b * (mse - mse_accum)
                mse_accum += delta / n

                report = (n - b) // 100 < n // 100

                optim.step()
                optim.zero_grad()
                model.clip()

                if report:
                    tokens = [
                        epoch + 1,
                        num_epochs,
                        n / N,
                        loss_accum,
                        acc_accum,
                        mse_accum,
                    ]
                    if output is not sys.stdout:
                        print(batch_report_fmt.format(*tokens), file=output)
                        output.flush()

            if (epoch + 1) % report_steps == 0:
                model.eval()

                with torch.no_grad():

                    (
                        inter_loss,
                        inter_correct,
                        inter_mse,
                        inter_pr,
                        inter_re,
                        inter_f1,
                        inter_aupr,
                    ) = interaction_eval(model, pairs_test_iterator, tensors, use_cuda)
                    tokens = [
                        epoch + 1,
                        num_epochs,
                        inter_loss,
                        inter_correct / (len(pairs_test_iterator) * batch_size),
                        inter_mse,
                        inter_pr,
                        inter_re,
                        inter_f1,
                        inter_aupr,
                    ]
                    print(epoch_report_fmt.format(*tokens), file=output)
                    output.flush()

                # Save the model
                if save_prefix is not None:
                    save_path = save_prefix + "_epoch" + str(epoch + 1).zfill(digits) + ".sav"
                    print(f"# Saving model to {save_path}", file=output)
                    model.cpu()
                    torch.save(model, save_path)
                    if use_cuda:
                        model.cuda()

            output.flush()

        if save_prefix is not None:
            save_path = save_prefix + "_final.sav"
            print(f"# Saving final model to {save_path}", file=output)
            model.cpu()
            torch.save(model, save_path)
            if use_cuda:
                model.cuda()

        output.close()

    if __name__ == "__main__":
        parser = argparse.ArgumentParser(description=__doc__)
        add_args(parser)
        main(parser.parse_args())