train.py

import pandas as pd
import torch
from tqdm import tqdm
import numpy as np
import json
import os
import wandb
from models.trainer import Trainer
from models.dream_models import DREAM_RNN, DREAM_CNN, DREAM_ATTN
from models.dl_utils import SeqExprDataset

#Defining constants
CUDA_DEVICE_ID = 0

HepG2 = f"human_mpra/HepG2_clean.tsv"
K562= f"human_mpra/K562_clean.tsv"
WTC11 = f"human_mpra/WTC11_clean.tsv"

MODEL_LOG_DIR = f"model_weights"
TRAIN_BATCH_SIZE = 32
N_PROCS = 4
VALID_BATCH_SIZE = 32
lr = 0.005 # 0.001 for DREAM-Attn, 0.005 for DREAM-CNN and DREAM-RNN
SEQ_SIZE = 230
NUM_EPOCHS = 80 #80
generator = torch.Generator()
generator.manual_seed(42)
device = torch.device(f"cuda:{CUDA_DEVICE_ID}")

#------------------Process Task ID---------------------------#
task_id = int(os.environ.get('SLURM_ARRAY_TASK_ID', 0)) 
array_id = task_id%10
training_id = task_id//10
model_id = training_id//3
type_id = training_id%3

type_pair = {0:"HepG2", 1:"K562", 2:"WTC11"}


TEST_ID = array_id
VALID_ID = array_id+1 if array_id<9 else 0

if(model_id==0):
    model_name = "rnn"
elif (model_id==1):
    model_name = "cnn"
elif (model_id==2):
    model_name = "attn"


#------------------Set up----------------------#
num_tasks = 10
seed = 42

# for human
seqsize=230
species = "human"
in_channels=5

#------------------ data--------------------------#

type_pair = {0:HepG2, 1:K562, 2:WTC11}
type_df = pd.read_csv(type_pair.get(type_id), sep='\t')




type_test = type_df[
    (type_df['fold'] == TEST_ID) 
].copy()

type_valid = type_df[
    (type_df['fold'] == VALID_ID) 
].copy()

type_train = type_df[
    (type_df['fold'] != TEST_ID)&(type_df['fold'] != VALID_ID) 
].copy()

rev_key = '_R' if type_id==2 else '_Reversed:'
type_test ['rev'] = type_test ['seq_id'].str.contains(rev_key, na=False).astype(int)
type_train['rev'] =  type_train['seq_id'].str.contains(rev_key, na=False).astype(int)
type_valid['rev'] = type_valid['seq_id'].str.contains(rev_key, na=False).astype(int)


BATCH_PER_EPOCH = len(type_train)//TRAIN_BATCH_SIZE
BATCH_PER_VALIDATION = len(type_valid)//TRAIN_BATCH_SIZE
#---------------- Train with model ---------------------#

ATTmodel = DREAM_ATTN(seqsize=seqsize,in_channels=in_channels)
RNNmodel = DREAM_RNN(seqsize=seqsize,in_channels=in_channels)
CNNmodel = DREAM_CNN(seqsize=seqsize,in_channels=in_channels)
model_dict = {
    "attn": ATTmodel,
    "rnn": RNNmodel,
    "cnn": CNNmodel
}
model = model_dict.get(model_name)
model_dir = f"{model_name}/{type_pair.get(type_id)}_{array_id}"
#train model
use_single_channel = species == 'yeast'
train_dataset = SeqExprDataset(df=type_train, 
                         species=species,
                         seqsize=seqsize,
                         use_single_channel=use_single_channel)
    
valid_dataset = SeqExprDataset(df=type_valid, 
                         species=species,
                         seqsize=seqsize,
                         use_single_channel=use_single_channel)
train_dl = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=TRAIN_BATCH_SIZE,
    shuffle=False,
    num_workers=4,
    generator=generator
)


valid_dl = torch.utils.data.DataLoader(
    valid_dataset,
    batch_size=VALID_BATCH_SIZE,
    shuffle=False,
    num_workers=4,
    generator=generator
)


trainer = Trainer(
        model=model,
        train_dataloader=train_dl,
        val_dataloader=valid_dl,
        model_dir=model_dir,
        num_epochs=80,
        lr=0.001 if model_name == "attn" else 0.005,
        device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
)
trainer.fit()