-
Notifications
You must be signed in to change notification settings - Fork 2
/
load_train.py
80 lines (61 loc) · 2.57 KB
/
load_train.py
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
from scipy.stats import spearmanr, pearsonr
from tqdm import tqdm
import numpy as np
import logging
import torch
def train_oiqa(epoch, net, criterion, optimizer, train_loader):
losses = []
net.train()
# save data for one epoch
pred_epoch = []
labels_epoch = []
for data in tqdm(train_loader):
d = data['d_img_org'].cuda()
labels = data['score']
name = data['name']
labels = torch.squeeze(labels.type(torch.FloatTensor)).cuda()
pred_d = net(d)
optimizer.zero_grad()
loss = criterion(torch.squeeze(pred_d), labels)
losses.append(loss.item())
loss.backward()
optimizer.step()
# save results in one epoch
pred_batch_numpy = pred_d.data.cpu().numpy()
labels_batch_numpy = labels.data.cpu().numpy()
pred_epoch = np.append(pred_epoch, pred_batch_numpy)
labels_epoch = np.append(labels_epoch, labels_batch_numpy)
# compute correlation coefficient
rho_s, _ = spearmanr(np.squeeze(pred_epoch), np.squeeze(labels_epoch))
rho_p, _ = pearsonr(np.squeeze(pred_epoch), np.squeeze(labels_epoch))
ret_loss = np.mean(losses)
logging.info('train epoch:{} / loss:{:.4} / SRCC:{:.4} / PLCC:{:.4}'.format(
epoch + 1, ret_loss, rho_s, rho_p))
return ret_loss, rho_s, rho_p
def eval_oiqa(config, epoch, net, criterion, test_loader):
with torch.no_grad():
losses = []
net.eval()
# save data for one epoch
pred_epoch = []
labels_epoch = []
for data in tqdm(test_loader):
d = data['d_img_org'].cuda()
labels = data['score']
name = data['name']
labels = torch.squeeze(labels.type(torch.FloatTensor)).cuda()
pred_d = net(d)
# compute loss
loss = criterion(torch.squeeze(pred_d), labels)
losses.append(loss.item())
# save results in one epoch
pred_batch_numpy = pred_d.data.cpu().numpy()
labels_batch_numpy = labels.data.cpu().numpy()
pred_epoch = np.append(pred_epoch, pred_batch_numpy)
labels_epoch = np.append(labels_epoch, labels_batch_numpy)
# compute correlation coefficient
rho_s, _ = spearmanr(np.squeeze(pred_epoch), np.squeeze(labels_epoch))
rho_p, _ = pearsonr(np.squeeze(pred_epoch), np.squeeze(labels_epoch))
logging.info('Epoch:{} ===== loss:{:.4} ===== SRCC:{:.4} ===== PLCC:{:.4}'.format(
epoch + 1, np.mean(losses), rho_s, rho_p))
return np.mean(losses), rho_s, rho_p