compara-deep-learning/train.py

import pickle
import numpy as np
import sys
import tensorflow as tf
from model import create_model


def create_branch_length_padding(bl):
    maxlen = 0
    for x in bl:
        if len(x) > maxlen:
            maxlen = len(x)

    for x in range(len(bl)):
        temp = bl[x]
        for i in range(len(temp), maxlen):
            temp = np.append(temp, [0])
        bl[x] = temp
    return bl


def train(train_synteny_matrices_global, train_synteny_matrices_local,
          train_pfam_matrices, train_branch_length_species,
          train_branch_length_homology_species, train_dist_p_s,
          train_dist_p_hs, train_distance, train_labels, v,
          num_epochs, learning_rate, decay, size_train,
          batch_size, model_name):
    print("Going to train model {} for:\n Batch Size:{} \
        \n Learning Rate:{} \n Decay:{}\n On {} Samples".format(
        v, batch_size, learning_rate, decay,
        len(train_synteny_matrices_global)
    ))
    graph, saver = create_model()
    synmg, synml, pfam, bls, blhs, dps, \
        dphs, dis, lr, y = graph.get_collection("input_nodes")
    loss, t_op, accuracy, init, summary = graph.get_collection("output_nodes")
    with tf.Session(graph=graph) as sess:
        writer = tf.summary.FileWriter('./'+model_name+'_v'+str(v), sess.graph)
        sess.run(init)
        learn = learning_rate
        for j in range(num_epochs):
            for i in range(size_train//batch_size):
                feed_dict = {
                    synmg: train_synteny_matrices_global
                    [i*batch_size:(i+1)*batch_size],
                    synml: train_synteny_matrices_local
                    [i*batch_size:(i+1)*batch_size],
                    pfam: train_pfam_matrices
                    [i*batch_size:(i+1)*batch_size]
                    .reshape((batch_size, 7, 7, 1)),
                    bls: train_branch_length_species
                    [i*batch_size:(i+1)*batch_size],
                    blhs: train_branch_length_homology_species
                    [i*batch_size:(i+1)*batch_size],
                    dps: train_dist_p_s
                    [i*batch_size:(i+1)*batch_size]
                    .reshape((batch_size, 1)),
                    dphs: train_dist_p_hs
                    [i*batch_size:(i+1)*batch_size]
                    .reshape((batch_size, 1)),
                    dis: train_distance
                    [i*batch_size:(i+1)*batch_size]
                    .reshape((batch_size, 1)),
                    lr: learn,
                    y: train_labels[i*batch_size:(i+1)*batch_size]
                }
                sess.run(t_op, feed_dict=feed_dict)

            test_dict = {
                synmg: train_synteny_matrices_global[size_train:],
                synml: train_synteny_matrices_local[size_train:],
                pfam: train_pfam_matrices[size_train:]
                .reshape((
                    len(train_synteny_matrices_global)-size_train, 7, 7, 1)),
                bls: train_branch_length_species[size_train:],
                blhs: train_branch_length_homology_species[size_train:],
                dps: train_dist_p_s[size_train:]
                .reshape((len(train_synteny_matrices_global)-size_train, 1)),
                dphs: train_dist_p_hs[size_train:]
                .reshape((len(train_synteny_matrices_global)-size_train, 1)),
                dis: train_distance[size_train:]
                .reshape((len(train_synteny_matrices_global)-size_train, 1)),
                y: train_labels[size_train:],
                lr: learn
            }
            accuracy_test, loss_test, summary_write = sess.run(
                [accuracy, loss, summary], feed_dict=test_dict)
            writer.add_summary(summary_write, i+1)
            print("Epoch:{} Test Accuracy:{} Test Loss:{}".format(
                j+1, accuracy_test*100, loss_test))
            learn *= decay
        saver.save(sess, model_name+"_v"+str(v)+"/model.ckpt")
        writer.close()


def read_positive(len_p, bls, blhs, dis, dps, dphs, sml, smg, pfam, label):
    rowsh = []
    with open("dataset", "rb") as file:
        rowsh = pickle.load(file)
    shi = np.random.permutation(len(rowsh))
    rows_shuffled = []
    for i in range(len(shi)):
        rows_shuffled.append(rowsh[shi[i]])
    rowsh = rows_shuffled
    spco = {}
    spcp = {}
    for row in rowsh:
        if row["species"] not in spco:
            spco[row["species"]] = 0
            spcp[row["species"]] = 0

    maxcount_o = int((len_p)*0.3/14)
    maxcount_p = int((len_p)*0.7/14)
    for row in rowsh:
        if row["label"] == 2:
            continue
        if row["label"] == 1 and spco[row["species"]] > maxcount_o:
            continue
        if row["label"] == 0 and spcp[row["species"]] > maxcount_p:
            continue
        bls.append(np.array(row["bls"]))
        blhs.append(np.array(row["blhs"]))
        dis.append(row["dis"])
        dps.append(row["dps"])
        dphs.append(row["dphs"])
        sml.append(row["local_alignment_matrix"])
        smg.append(row["global_alignment_matrix"])
        pfam.append(row["pfam_matrix"])
        label.append(row["label"])
        if row["label"] == 1:
            spco[row["species"]] += 1
        if row["label"] == 0:
            spcp[row["species"]] += 1


def read_negative(len_n, bls, blhs, dis, dps, dphs, sml, smg, pfam, label):
    rows = []
    with open("dataset", "rb") as file:
        rows = pickle.load(file)
    rows = [row for row in rows if row["label"] == 2]
    shi = np.random.permutation(len(rows))
    rows_shuffled = []
    for i in range(len(shi)):
        rows_shuffled.append(rows[shi[i]])
    rows = rows_shuffled
    rows = rows[:len_n]
    for row in rows:
        bls.append(np.array(row["bls"]))
        blhs.append(np.array(row["blhs"]))
        dis.append(row["dis"])
        dps.append(row["dps"])
        dphs.append(row["dphs"])
        sml.append(row["local_alignment_matrix"])
        smg.append(row["global_alignment_matrix"])
        label.append(row["label"])
        pfam.append(row["pfam_matrix"])


def train_models(model_name, start, end, num_epochs, learn_rate,
                 decay, size_train, batch_size):
    k = 1
    for i in range(start//10, end//10+1):
        bls = []
        blhs = []
        dis = []
        dps = []
        dphs = []
        sml = []
        smg = []
        pfam = []
        label = []
        portion = float(i/10)
        len_n = int(size_train*portion)
        len_p = int(size_train*(1-portion))
        read_positive(len_p, bls, blhs, dis, dps, dphs, sml, smg, pfam, label)
        read_negative(len_n, bls, blhs, dis, dps, dphs, sml, smg, pfam, label)
        bls = create_branch_length_padding(bls)
        blhs = create_branch_length_padding(blhs)
        bls = np.array(bls)
        print(bls.shape)
        blhs = np.array(blhs)
        print(blhs.shape)
        dis = np.array(dis)
        print(dis.shape)
        dps = np.array(dps)
        print(dps.shape)
        dphs = np.array(dphs)
        print(dphs.shape)
        sml = np.array(sml)
        print(sml.shape)
        smg = np.array(smg)
        print(smg.shape)
        pfam = np.array(pfam)
        print(pfam.shape)
        label = np.array(label)
        print(label.shape)
        shi = np.random.permutation(len(label))
        labels = label[shi]
        bls = bls[shi]
        blhs = blhs[shi]
        dis = dis[shi]
        dps = dps[shi]
        dphs = dphs[shi]
        sml = sml[shi]
        smg = smg[shi]
        pfam = pfam[shi]
        train(smg, sml, pfam, bls, blhs, dps, dphs, dis, labels, k, num_epochs,
              learn_rate, decay, int(0.9*size_train), batch_size, model_name)
        k += 1


def main():
    arg = sys.argv
    model_name = arg[-8]
    start_p = int(arg[-7])
    end_p = int(arg[-6])
    num_epochs = int(arg[-5])
    learn_rate = float(arg[-4])
    decay = float(arg[-3])
    size_train = float(arg[-2])
    batch_size = int(arg[-1])
    train_models(model_name, start_p, end_p, num_epochs,
                 learn_rate, decay, size_train, batch_size)


if __name__ == "__main__":
    main()