Train

In [ ]:

import torch
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
from rouge import Rouge
import tqdm
import time
pad_token = 0
bos_token = 1
eos_token = 2
def train(model, loss_func, optimizer, data_loader, device, epochs):
    start_time = time.time()
    model = model.to(device)
    for epoch in range(epochs):
        model.train()
        total_loss = 0.0
        for _, (description, diagnosis, description_len, diagnosis_len) in tqdm.tqdm(enumerate(data_loader), total=len(data_loader)):
            y_in=diagnosis[:, :-1]
            y_target=diagnosis[:, 1:]
            # padding
            # 开头和末尾分别添加标记
            y_in = torch.cat((torch.tensor(bos_token).repeat(diagnosis.shape[0],1),y_in), dim=1)
            y_in = y_in.to(device)
            y_target = torch.cat((y_target, torch.tensor(pad_token).repeat(diagnosis.shape[0], 1)), dim=1)
            y_target = y_target.to(device)

            # 提取description和diagnosis信息
            description = description.to(device)
            diagnosis = diagnosis.to(device)
            description_len = description_len.to(device)
            diagnosis_len = diagnosis_len.to(device)
            #add
            optimizer.zero_grad()
            y_pred, _ = model(description, y_in)
            y_pred = y_pred.permute(0,2,1)
            loss = loss_func(y_pred, y_target)
            loss.backward()
            optimizer.step()
            total_loss += loss.item() / y_target.shape[0]

        total_loss /= len(data_loader)
        print(f'Epoch {epoch} train loss: {total_loss}')
    
    end_time = time.time()
    print(f'Train cost: {end_time - start_time} s')

import torch from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction from rouge import Rouge import tqdm import time pad_token = 0 bos_token = 1 eos_token = 2 def train(model, loss_func, optimizer, data_loader, device, epochs): start_time = time.time() model = model.to(device) for epoch in range(epochs): model.train() total_loss = 0.0 for _, (description, diagnosis, description_len, diagnosis_len) in tqdm.tqdm(enumerate(data_loader), total=len(data_loader)): y_in=diagnosis[:, :-1] y_target=diagnosis[:, 1:] # padding # 开头和末尾分别添加标记 y_in = torch.cat((torch.tensor(bos_token).repeat(diagnosis.shape[0],1),y_in), dim=1) y_in = y_in.to(device) y_target = torch.cat((y_target, torch.tensor(pad_token).repeat(diagnosis.shape[0], 1)), dim=1) y_target = y_target.to(device) # 提取description和diagnosis信息 description = description.to(device) diagnosis = diagnosis.to(device) description_len = description_len.to(device) diagnosis_len = diagnosis_len.to(device) #add optimizer.zero_grad() y_pred, _ = model(description, y_in) y_pred = y_pred.permute(0,2,1) loss = loss_func(y_pred, y_target) loss.backward() optimizer.step() total_loss += loss.item() / y_target.shape[0] total_loss /= len(data_loader) print(f'Epoch {epoch} train loss: {total_loss}') end_time = time.time() print(f'Train cost: {end_time - start_time} s')

In [ ]:

def rouge(preds, labels):
    ans = 0
    rouge = Rouge()
    for i in range(len(preds)):
        if hasattr(labels[i], '__len__'):
            ref= ' '.join([str(num) for num in labels[i]])
        else:
            ref=str(labels[i])
        tmp= ' '.join([str(num) for num in preds[i]])
        ans += rouge.get_scores(tmp, ref)[0]['rouge-2']['r']
    ans /= len(preds)
    return ans

def rouge(preds, labels): ans = 0 rouge = Rouge() for i in range(len(preds)): if hasattr(labels[i], '__len__'): ref= ' '.join([str(num) for num in labels[i]]) else: ref=str(labels[i]) tmp= ' '.join([str(num) for num in preds[i]]) ans += rouge.get_scores(tmp, ref)[0]['rouge-2']['r'] ans /= len(preds) return ans

In [ ]:

def bleu(preds, labels):
    ans = 0
    for i in range(len(preds)):
        if hasattr(labels[i], '__len__'):
            ref= ' '.join([str(num) for num in labels[i]])
        else:
            ref=str(labels[i])
        tmp=' '.join([str(num) for num in preds[i]])
        ans+=sentence_bleu([ref], tmp, smoothing_function=SmoothingFunction().method4)
    ans/=len(preds)
    return ans

def bleu(preds, labels): ans = 0 for i in range(len(preds)): if hasattr(labels[i], '__len__'): ref= ' '.join([str(num) for num in labels[i]]) else: ref=str(labels[i]) tmp=' '.join([str(num) for num in preds[i]]) ans+=sentence_bleu([ref], tmp, smoothing_function=SmoothingFunction().method4) ans/=len(preds) return ans

In [ ]:

def valid(model, data_loader, device, max_len = 160):
    model.eval()
    preds = []
    labels = []
    model = model.to(device)
    for i, (description, diagnosis, description_len, diagnosis_len) in tqdm.tqdm(enumerate(data_loader), total=len(data_loader)):
        description = description.to(device)
        diagnosis = diagnosis.to(device)
        description_len = description_len.to(device)
        diagnosis_len = diagnosis_len.to(device)

        batches = description.size(0)
        for j in range(batches):
            description_sequence = description[j, :description_len[j]].unsqueeze(0)
            diagnosis_sequence = diagnosis[j, 1:diagnosis_len[j]].cpu().unsqueeze(0)
            pred = model.predict(description_sequence, description_len[j], max_len, device)
            preds.append(pred)
            labels.append(diagnosis_sequence.squeeze().tolist())

    
    bleu_score = bleu(preds, labels)
    rouge_score = rouge(preds, labels)
    
    print(f"BLEU: {bleu_score}")
    print(f"ROUGE: {rouge_score}")      

def valid(model, data_loader, device, max_len = 160): model.eval() preds = [] labels = [] model = model.to(device) for i, (description, diagnosis, description_len, diagnosis_len) in tqdm.tqdm(enumerate(data_loader), total=len(data_loader)): description = description.to(device) diagnosis = diagnosis.to(device) description_len = description_len.to(device) diagnosis_len = diagnosis_len.to(device) batches = description.size(0) for j in range(batches): description_sequence = description[j, :description_len[j]].unsqueeze(0) diagnosis_sequence = diagnosis[j, 1:diagnosis_len[j]].cpu().unsqueeze(0) pred = model.predict(description_sequence, description_len[j], max_len, device) preds.append(pred) labels.append(diagnosis_sequence.squeeze().tolist()) bleu_score = bleu(preds, labels) rouge_score = rouge(preds, labels) print(f"BLEU: {bleu_score}") print(f"ROUGE: {rouge_score}")