modified_run.py

from torch import nn
import torch
import pandas as pd
import pickle, time
import os
import torch.nn.functional as F
import model, parsing, utils, hf
import torch.optim as optim
from tqdm import tqdm
from torch.utils.data import DataLoader
from torch.profiler import profile, record_function, ProfilerActivity
import numpy as np
import torch.autograd.profiler as profiler
import torch
import torch.nn
import torch.optim


def train(model, train_dataset):
    print("Starting training ........")
    dataloader = DataLoader(train_dataset, batch_size=20)
    train_loss = 0.0
    batch_count = 0
    model.train()
    loading = time.time()

    for batch in dataloader:
        print("data loader time :", time.time() - loading)  # 현재시각 - 시작시간 = 실행 시간
        measure = time.time()
        optimizer.zero_grad()
        # if batch_count % 500 == 0:
        #    print(f"Starting batch: {batch_count}")
        batch_count += 1
        if batch_count == 4:
            break
        print(f"Starting batch: {batch_count}")
        context, question, start, end = batch

        context, question, start, end = (
            context.to(device),
            question.to(device),
            start.to(device),
            end.to(device),
        )
        """
        from torchsummaryX import summary
        summary(model, context, question)
        preds = model(context, question)  # warm-up
        """
        """
        with profiler.profile(
            with_stack=True, use_cuda=False, profile_memory=True
        ) as prof:
            preds = model(context, question)

        print(
            prof.key_averages(group_by_input_shape=True).table(
                sort_by="cuda_memory_usage", row_limit=10
            )
        )
        """
        preds = model(context, question)
        start_pred, end_pred = preds  # p1, p2
        print("model computing time :", time.time() - measure)  # 현재시각 - 시작시간 = 실행 시간
        s_idx, e_idx = start, end
        measure = time.time()

        loss = F.cross_entropy(start_pred, s_idx) + F.cross_entropy(end_pred, e_idx)

        loss.backward()

        optimizer.step()

        train_loss += loss.item()
        print("loss computing time :", time.time() - measure)  # 현재시각 - 시작시간 = 실행 시간
        loading = time.time()

    return train_loss / len(train_dataset)


def valid(model, valid_dataset):

    print("Starting validation .........")

    valid_loss = 0.0

    batch_count = 0

    f1, em = 0.0, 0.0

    model.eval()

    predictions = {}

    for batch in valid_dataset:

        if batch_count % 500 == 0:
            print(f"Starting batch {batch_count}")
        batch_count += 1

        context, question, char_ctx, char_ques, label, ctx, answers, ids = batch
        context, question, char_ctx, char_ques, label = (
            context.to(device),
            question.to(device),
            char_ctx.to(device),
            char_ques.to(device),
            label.to(device),
        )

        with torch.no_grad():

            s_idx, e_idx = label[:, 0], label[:, 1]

            preds = model(context, question, char_ctx, char_ques)

            p1, p2 = preds

            loss = F.cross_entropy(p1, s_idx) + F.cross_entropy(p2, e_idx)

            valid_loss += loss.item()

            batch_size, c_len = p1.size()

            st = nn.Softmax(dim=1)
            # 근데 stable한거는 알겠는데 softmax와 logsoftmax 차이는 분명 존재하는데;
            mask = (
                (torch.ones(c_len, c_len) * float("-inf"))
                .to(device)
                .tril(-1)
                .unsqueeze(0)
                .expand(batch_size, -1, -1)
            )
            score = torch.bmm(st(p1).unsqueeze(2), st(p2).unsqueeze(1)) + mask
            score, s_idx = score.max(dim=1)
            score, e_idx = score.max(dim=1)
            s_idx = torch.gather(s_idx, 1, e_idx.view(-1, 1)).squeeze()

            for i in range(batch_size):
                id = ids[i]
                pred = context[i][s_idx[i] : e_idx[i] + 1]
                pred = " ".join([idx2word[idx.item()] for idx in pred])
                predictions[id] = pred

    em, f1 = utils.evaluate(predictions)
    return valid_loss / len(valid_dataset), em, f1


#######  MAIN  #######

# load data from pickle files

file1 = os.path.exists("parsing/bidaftrain.pkl")
file2 = os.path.exists("parsing/bidafvalid.pkl")
file3 = os.path.exists("parsing/qanetw2id.pickle")
file4 = os.path.exists("parsing/qanetc2id.pickle")
glove = os.path.exists("parsing/bidafglove.npy")
isFile = file1 and file2 and file3 and file4 and glove

if not isFile:
    print("start getting files")
    parsing.get_datafiles()

train_df = pd.read_pickle("parsing/bidaftrain.pkl")
valid_df = pd.read_pickle("parsing/bidafvalid.pkl")

with open("parsing/qanetw2id.pickle", "rb") as handle:
    word2idx = pickle.load(handle)
with open("parsing/qanetc2id.pickle", "rb") as handle:
    char2idx = pickle.load(handle)

idx2word = {v: k for k, v in word2idx.items()}


device = torch.device("cuda:3" if torch.cuda.is_available else "cpu")
#device = torch.device("cpu")

train_dataset = hf.SQUAD()
valid_dataset = hf.SQUAD()


# below are codes for checking if this model causes errors
"""
k= train_df[:2000]
z= valid_df[:700]
train_dataset = squad.SquadDataset(k,16)
valid_dataset = squad.SquadDataset(z,16)
"""


CHAR_VOCAB_DIM = len(char2idx)
EMB_DIM = 100
CHAR_EMB_DIM = 8
NUM_OUTPUT_CHANNELS = 100
KERNEL_SIZE = (8, 5)
HIDDEN_DIM = 100


myModel = model.BIDAF(
    EMB_DIM, KERNEL_SIZE, CHAR_VOCAB_DIM, CHAR_EMB_DIM, HIDDEN_DIM, device
).to(device)
optimizer = optim.Adadelta(myModel.parameters(), lr=0.5)
# myModel = nn.DataParallel(myModel, device_ids=[0,1])


train_losses = []
valid_losses = []
ems = []
f1s = []
epochs = 12


for epoch in tqdm(range(epochs)):
    print(f"Epoch {epoch+1}")
    start_time = time.time()

    train_loss = train(myModel, train_dataset)
    break
    # valid_loss, em, f1 = valid(myModel, valid_dataset)
"""
    torch.save(
        {
            "epoch": epoch,
            "model_state_dict": myModel.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "loss": valid_loss,
            "em": em,
            "f1": f1,
        },
        "result/bidaf_run4_{}.pth".format(epoch),
    )

    end_time = time.time()

    epoch_mins, epoch_secs = utils.epoch_time(start_time, end_time)

    train_losses.append(train_loss)
    valid_losses.append(valid_loss)
    ems.append(em)
    f1s.append(f1)

    print(f"Epoch train loss : {train_loss}| Time: {epoch_mins}m {epoch_secs}s")
    print(f"Epoch valid loss: {valid_loss}")
    print(f"Epoch EM: {em}")
    print(f"Epoch F1: {f1}")
    print(
        "===================================================================================="
        
    )
"""