diffqkt/run.py at master · ddsplus/diffqkt · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
from sklearn import metrics
from tqdm import tqdm
import torch
import numpy as np
import torch.nn as nn
import time

from load_data import getLoader

def run_epoch(max_problem, pro_path, skill_path, batch_size, is_train, min_problem_num, max_problem_num,
              model, optimizer, criterion, device, grad_clip):
    total_correct = 0
    total_num = 0
    total_loss = []

    labels = []
    outputs = []

    train_time = 0.0
    test_time = 0.0

    if is_train:
        model.train()
    else:
        model.eval()

    data_loader = getLoader(max_problem, pro_path, skill_path, batch_size, is_train, min_problem_num, max_problem_num)

    for i, data in tqdm(enumerate(data_loader), desc='加载中...'):

        last_problem, last_skill, last_ans, next_problem, next_skill, next_ans, mask = data

        if is_train:
            step_start_time = time.time()

            predict, _, f_loss = model(last_problem, last_skill, last_ans, next_problem, next_skill, next_ans, mask)

            next_predict = torch.masked_select(predict, mask)
            next_true = torch.masked_select(next_ans, mask)

            loss = criterion(next_predict, next_true) + f_loss

            optimizer.zero_grad()
            loss.backward()

            nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            optimizer.step()

            step_end_time = time.time()
            step_time = step_end_time - step_start_time
            train_time += step_time

            labels.extend(next_true.view(-1).data.cpu().numpy())
            outputs.extend(next_predict.view(-1).data.cpu().numpy())

            total_loss.append(loss.item())
            total_num += len(next_true)

            to_pred = (next_predict >= 0.5).long()
            total_correct += (next_true == to_pred).sum()

            print(f"Batch {i+1} training completed in {step_time:.2f}s")
        else:
            step_start_time = time.time()

            with torch.no_grad():
                predict, _, _ = model(last_problem, last_skill, last_ans, next_problem, next_skill, next_ans, mask, False)

                next_predict = torch.masked_select(predict, mask)
                next_true = torch.masked_select(next_ans, mask)

                loss = criterion(next_predict, next_true)

                labels.extend(next_true.view(-1).data.cpu().numpy())
                outputs.extend(next_predict.view(-1).data.cpu().numpy())

                total_loss.append(loss.item())
                total_num += len(next_true)

                to_pred = (next_predict >= 0.5).long()
                total_correct += (next_true == to_pred).sum()

            step_end_time = time.time()
            step_time = step_end_time - step_start_time
            test_time += step_time

            print(f"Batch {i+1} test completed in {step_time:.2f}s")

    avg_loss = np.average(total_loss)
    acc = total_correct * 1.0 / total_num
    auc = metrics.roc_auc_score(labels, outputs)

    if is_train:
        print(f"\n总训练时间: {train_time:.2f}s")
    else:
        print(f"\n总测试时间: {test_time:.2f}s")

    return avg_loss, acc, auc