forked from hengyuan-hu/bottom-up-attention-vqa
-
Notifications
You must be signed in to change notification settings - Fork 5
/
train.py
116 lines (89 loc) · 3.42 KB
/
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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
import json
import os
import pickle
import time
from os.path import join
import torch
import torch.nn as nn
import utils
from torch.autograd import Variable
import numpy as np
from tqdm import tqdm
def compute_score_with_logits(logits, labels):
logits = torch.max(logits, 1)[1].data # argmax
one_hots = torch.zeros(*labels.size()).cuda()
one_hots.scatter_(1, logits.view(-1, 1), 1)
scores = (one_hots * labels)
return scores
def train(model, train_loader, eval_loader, num_epochs, output, eval_each_epoch):
utils.create_dir(output)
optim = torch.optim.Adamax(model.parameters())
logger = utils.Logger(os.path.join(output, 'log.txt'))
all_results = []
total_step = 0
for epoch in range(num_epochs):
total_loss = 0
train_score = 0
t = time.time()
for i, (v, q, a, b) in tqdm(enumerate(train_loader), ncols=100,
desc="Epoch %d" % (epoch+1), total=len(train_loader)):
total_step += 1
v = Variable(v).cuda()
q = Variable(q).cuda()
a = Variable(a).cuda()
b = Variable(b).cuda()
pred, loss = model(v, None, q, a, b)
if (loss != loss).any():
raise ValueError("NaN loss")
loss.backward()
nn.utils.clip_grad_norm(model.parameters(), 0.25)
optim.step()
optim.zero_grad()
batch_score = compute_score_with_logits(pred, a.data).sum()
total_loss += loss.data[0] * v.size(0)
train_score += batch_score
total_loss /= len(train_loader.dataset)
train_score = 100 * train_score / len(train_loader.dataset)
run_eval = eval_each_epoch or (epoch == num_epochs - 1)
if run_eval:
model.train(False)
results = evaluate(model, eval_loader)
results["epoch"] = epoch+1
results["step"] = total_step
results["train_loss"] = total_loss
results["train_score"] = train_score
all_results.append(results)
with open(join(output, "results.json"), "w") as f:
json.dump(all_results, f, indent=2)
model.train(True)
eval_score = results["score"]
bound = results["upper_bound"]
logger.write('epoch %d, time: %.2f' % (epoch+1, time.time()-t))
logger.write('\ttrain_loss: %.2f, score: %.2f' % (total_loss, train_score))
if run_eval:
logger.write('\teval score: %.2f (%.2f)' % (100 * eval_score, 100 * bound))
model_path = os.path.join(output, 'model.pth')
torch.save(model.state_dict(), model_path)
def evaluate(model, dataloader):
score = 0
upper_bound = 0
num_data = 0
all_logits = []
all_bias = []
for v, q, a, b in tqdm(dataloader, ncols=100, total=len(dataloader), desc="eval"):
v = Variable(v, volatile=True).cuda()
q = Variable(q, volatile=True).cuda()
pred, _ = model(v, None, q, None, None)
all_logits.append(pred.data.cpu().numpy())
batch_score = compute_score_with_logits(pred, a.cuda()).sum()
score += batch_score
upper_bound += (a.max(1)[0]).sum()
num_data += pred.size(0)
all_bias.append(b)
score = score / len(dataloader.dataset)
upper_bound = upper_bound / len(dataloader.dataset)
results = dict(
score=score,
upper_bound=upper_bound,
)
return results