model/src/gan/wgan.py

import numpy as np
import pandas as pd
import pdb
import tensorflow as tf
import sys
# from IPython.display import clear_output
import time
from lib import models
from lib import utils
import socket
import datetime
from tqdm import tqdm
import os
import random
import matplotlib.pyplot as plt
import argparse

import tensorflow.keras.backend as K
from tensorflow.keras.optimizers import Adam

tf.compat.v1.enable_eager_execution()


MODEL_NAME = 'WGAN-TF2'
OUTPUT_PATH = os.path.join('outputs', MODEL_NAME)
TRAIN_LOGDIR = os.path.join("logs_", "tensorflow", MODEL_NAME, 'train_data') # Sets up a log directory.
if not os.path.exists(OUTPUT_PATH):
    os.makedirs(OUTPUT_PATH)



file_writer = tf.summary.create_file_writer(TRAIN_LOGDIR)

parser = argparse.ArgumentParser()
parser.add_argument('-d', type=str, required=False ,default='./../../data/utrdb2.csv')    
parser.add_argument('-bs', type=int, required=False ,default=64)
parser.add_argument('-lr', type=int, required=False ,default=5)
parser.add_argument('-mil', type=int, required=False ,default=64)
parser.add_argument('-mxl', type=int, required=False ,default=128)
parser.add_argument('-dim', type=int, required=False ,default=40)
parser.add_argument('-gpu', type=str, required=False ,default='-1')
args = parser.parse_args()

if args.gpu == '-1':
    device = 'cpu'
else:
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
    device = 'cuda'

def plot(x, y, logdir, name, xlabel=None, ylabel=None, title=None):

    plt.plot(x,y,'-')
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.savefig(logdir+'/'+name+'.png')
    plt.clf()


def plot_valid(x1, y1, x2, y2, logdir, name, xlabel=None, ylabel=None, title=None):

    plt.plot(x2,y2,'-',color='tab:blue')
    plt.plot(x1,y1,'-',color='tab:orange')

    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)

    plt.savefig(logdir+'/'+name+'.png')
    plt.clf()    

def gradient_penalty_loss( y_true, y_pred, discriminator):
  """
  Computes gradient penalty based on prediction and weighted real / fake samples
  """
  alpha = K.random_uniform((DIM, 1, 1))
  averaged_samples = (alpha * y_pred) + ((1 - alpha) * y_true)

  gradients = K.gradients(y_pred, averaged_samples)[0]
  # compute the euclidean norm by squaring ...
  gradients_sqr = K.square(gradients)
  #   ... summing over the rows ...
  gradients_sqr_sum = K.sum(gradients_sqr,
                            axis=np.arange(1, len(gradients_sqr.shape)))
  #   ... and sqrt
  gradient_l2_norm = K.sqrt(gradients_sqr_sum)
  # compute lambda * (1 - ||grad||)^2 still for each single sample
  gradient_penalty = K.square(1 - gradient_l2_norm)
  # return the mean as loss over all the batch samples
  return K.mean(gradient_penalty)

def log(samples_dir=False,suff=None):
    stamp = datetime.date.strftime(datetime.datetime.now(), "%Y.%m.%d-%Hh%Mm%Ss") + "_{}".format(socket.gethostname())
    full_logdir = os.path.join("./logs/", stamp)
    if suff:
        full_logdir = full_logdir + suff
    os.makedirs(full_logdir, exist_ok=True)
    if samples_dir: os.makedirs(os.path.join(full_logdir, "samples"), exist_ok=True)
    log_dir = "{}:{}".format(socket.gethostname(), full_logdir)
    
    return full_logdir, 0

data_path = args.d
data_utr = pd.read_csv(data_path)
UTRdf = data_utr['seq'].to_numpy()

seqs = []

UTR_LEN = args.mxl

for i in range(len(UTRdf)):
    if len(UTRdf[i]) < UTR_LEN+1 and len(UTRdf[i]) > int(UTR_LEN/2):
        if UTRdf[i] not in seqs:
            seqs.append(UTRdf[i])

sequences = np.array(seqs)

sequences = [x.upper() for x in sequences]

rna_vocab = {"A":0,
             "C":1,
             "G":2,
             "U":3,
             "*":4}

rev_rna_vocab = {v:k for k,v in rna_vocab.items()}

def one_hot_encode(seq, SEQ_LEN=UTR_LEN):
    mapping = dict(zip("ACGT*", range(5)))    
    seq2 = [mapping[i] for i in seq]
    if len(seq2) < SEQ_LEN:
        extra = [np.eye(5)[4]] * (SEQ_LEN - len(seq2))
        return np.vstack([np.eye(5)[seq2] , extra])
    return np.eye(5)[seq2]

def one_hot_encode_2(seq, SEQ_LEN=UTR_LEN):
    mapping = dict(zip("ACGT", range(4)))    
    seq2 = [mapping[i] for i in seq]
    return np.eye(4)[seq2]    

ohe_sequences = np.asarray([one_hot_encode(x) for x in sequences])


BATCH_SIZE = args.bs # Batch size
ITERS = 4000 # How many iterations to train for
SEQ_LEN = UTR_LEN # Sequence length in characters
DIM = args.dim # Model dimensionality.
CRITIC_ITERS = 5 # How many critic iterations per generator iteration. 
LAMBDA = 10 # Gradient penalty lambda hyperparameter.
LR = np.power(np.e,args.lr)


LAMBDA = 10 # For gradient penalty

CURRENT_EPOCH = 1 # Epoch start from
SAVE_EVERY_N_EPOCH = 50 # Save checkpoint at every n epoch

MIN_LR = 0.000001 # Minimum value of learning rate
DECAY_FACTOR=1.00004 # learning rate decay factor
'''
Set seed for reproducibility
'''
seed = 35
np.random.seed(seed)
# tf.set_random_seed(seed)
tf.random.set_seed(seed)

logdir, checkpoint_baseline = log(samples_dir=True)
logdir2 = ''

'''
Build GAN
'''
model_type = "resnet"
data_enc_dim = 5
data_size = SEQ_LEN * data_enc_dim
# data_size = 256
gen_layers = 3
disc_layers = 3
lmbda = 10. #lipschitz penalty hyperparameter.

SAMPLE_SIZE = 128

N_CHANNELS = DIM

G = models.resnet_g2(DIM,N_CHANNELS,SEQ_LEN,5,res_layers=gen_layers)
D = models.resnet_d2(N_CHANNELS,SEQ_LEN,5,res_layers=disc_layers)

G.summary()
D.summary()

D_optimizer = Adam(learning_rate=LR, beta_1=0.5, beta_2=0.99)
G_optimizer = Adam(learning_rate=LR, beta_1=0.5, beta_2=0.99)

EPOCHs = ITERS

@tf.function
def WGAN_GP_train_d_step(real_sample, batch_size, step):

    noise = tf.random.normal([batch_size, DIM])
    epsilon = tf.random.uniform(shape=[batch_size, 1, 1], minval=0, maxval=1)
    ###################################
    # Train D
    ###################################
    with tf.GradientTape(persistent=True) as d_tape:
        with tf.GradientTape() as gp_tape:
            fake_sample = G([noise], training=True)
            fake_sample_mixed = epsilon * tf.dtypes.cast(real_sample, tf.float32) + ((1 - epsilon) * fake_sample)
            fake_mixed_pred = D([fake_sample_mixed], training=True)
            
        # Compute gradient penalty
        grads = gp_tape.gradient(fake_mixed_pred, fake_sample_mixed)
        grad_norms = tf.sqrt(tf.reduce_sum(tf.square(grads), axis=[1,2])) # Originally axis=[1,2]
        # grad_norms = tf.norm(grads, axis=[1,2])
        gradient_penalty = tf.reduce_mean(tf.square(grad_norms - 1.))
        
        fake_pred = D([fake_sample], training=True)
        real_pred = D([real_sample], training=True)
        
        D_loss = tf.reduce_mean(fake_pred) - tf.reduce_mean(real_pred) + LAMBDA * gradient_penalty
    # Calculate the gradients for discriminator
    D_gradients = d_tape.gradient(D_loss,D.trainable_variables)
    # Apply the gradients to the optimizer
    D_optimizer.apply_gradients(zip(D_gradients,D.trainable_variables))
    # Write loss values to tensorboard
    if step % 10 == 0:
        with file_writer.as_default():
            tf.summary.scalar('D_loss', tf.reduce_mean(D_loss), step=step)

    return D_loss, gradient_penalty

@tf.function
def WGAN_GP_train_g_step(real_sample, batch_size, step):

    noise = tf.random.normal([batch_size, DIM])
    ###################################
    # Train G
    ###################################
    with tf.GradientTape() as g_tape:
        fake_sample = G([noise], training=True)
        fake_pred = D([fake_sample], training=True)
        G_loss = -tf.reduce_mean(fake_pred)

    G_gradients = g_tape.gradient(G_loss,
                                            G.trainable_variables)
    # Apply the gradients to the optimizer
    G_optimizer.apply_gradients(zip(G_gradients,
                                                G.trainable_variables))
    # Write loss values to tensorboard
    if step % 10 == 0:
        with file_writer.as_default():
            tf.summary.scalar('G_loss', G_loss, step=step)

    return G_loss, noise

checkpoint_path = os.path.join("checkpoints", "tensorflow", MODEL_NAME)

ckpt = tf.train.Checkpoint(generator=G,
                        discriminator=D,
                        G_optimizer=G_optimizer,
                        D_optimizer=D_optimizer)

ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=40)

def generate_and_save_samples(model, epoch, test_input, figure_size=(12,6), subplot=(3,6), save=True, is_flatten=False):
    '''
        Generate samples and plot it.
    '''
    predictions = model.predict(test_input)
    utils.save_samples(logdir, predictions, epoch, rev_rna_vocab, annotated=False)


'''
load data
'''

Train = True
validate = True

data = ohe_sequences


if validate:
    split = len(data) // 10
    # print(split)
    valid_data = data[:split]
    train_data = data[split:]
    if len(train_data) == 1: train_data = train_data[0]
    if len(valid_data) == 1: valid_data = valid_data[0]
else:
    train_data = data


train_data = train_data.astype('float32')
valid_data = valid_data.astype('float32')

train_seqs = tf.data.Dataset.from_tensor_slices(train_data).shuffle(len(train_data)).batch(BATCH_SIZE)
valid_seqs = tf.data.Dataset.from_tensor_slices(valid_data).shuffle(len(valid_data)).batch(split)


plot_iter = 10

current_learning_rate = LR
trace = True
n_critic_count = 0

d_losses = []
g_losses = []
gradient_penalties = []
iterations = 0
iteration_numbers = []
iteration_numbers_valid = []
d_validation_losses = []

random_name = time.strftime("%Y%m%d-%H%M%S")

gen_iters = 0

if Train:
    sample_noise = tf.random.normal([BATCH_SIZE, DIM])
    generate_and_save_samples(G, 0, [sample_noise], figure_size=(12,6), subplot=(3,6), save=False, is_flatten=False)

    pbar = tqdm(range(EPOCHs))
    for epoch in pbar:
        start = time.time()

        tdataset = train_seqs.enumerate()

        for step, tdata in tdataset.as_numpy_iterator():
            current_batch_size = tdata.shape[0]

            d_loss, gp = WGAN_GP_train_d_step(tdata, batch_size=tf.constant(current_batch_size, dtype=tf.int64), step=tf.constant(step, dtype=tf.int64))
            n_critic_count += 1
            if n_critic_count >= CRITIC_ITERS: 
                g_loss, noise = WGAN_GP_train_g_step(tdata, batch_size= tf.constant(current_batch_size, dtype=tf.int64), step=tf.constant(step, dtype=tf.int64))
                gen_iters += 1
                n_critic_count = 0
            
            if step % 10 == 0:
                print ('.', end='')

        
        if epoch % SAVE_EVERY_N_EPOCH == 0 and epoch != 0:
            ckpt_save_path = ckpt_manager.save()
            utils.save_checkpoints(logdir,G,epoch)

        if epoch % 50 == 0:    
            
            generate_and_save_samples(G, epoch, [sample_noise], figure_size=(12,6), subplot=(3,6), save=True, is_flatten=False)

        os.system('clear')

        iteration_numbers.append(iterations)
        g_losses.append(-g_loss)
        d_losses.append(-d_loss)
        gradient_penalties.append(gp)
        # plot(iteration_numbers, d_losses, logdir, 'discriminator_loss', xlabel="Iteration", ylabel="Discriminator Cost")
        # plot(iteration_numbers, g_losses, logdir, 'generator_loss', xlabel="Iteration", ylabel="Generator Cost")
        # plot(iteration_numbers, gradient_penalties, logdir, 'gradient_penalty', xlabel="Iteration", ylabel="Gradient Penalty")

        iterations+=1

        iteration_numbers_valid.append(iterations)
        fake_sample_valid = G([noise], training=True)
        fake_pred_valid = D([fake_sample_valid], training=True)
        real_pred_valid = D([tf.convert_to_tensor(list(valid_seqs.as_numpy_iterator())[0])], training=True)
        
        D_loss_valid = tf.reduce_mean(fake_pred_valid) - tf.reduce_mean(real_pred_valid)
        d_validation_losses.append(-D_loss_valid)

        # plot_valid(iteration_numbers_valid, d_validation_losses, iteration_numbers, d_losses, logdir, 'validation_loss', xlabel="Iteration", ylabel="D Validation Loss")

    ckpt_save_path = ckpt_manager.save()
    print ('Saving checkpoint for epoch {} at {}'.format(EPOCHs,
                                                            ckpt_save_path))


print("####################################################")
print(f"############### Gen Iterations : {gen_iters} #####")
print("####################################################")