-
Notifications
You must be signed in to change notification settings - Fork 3
/
train_test_functions.py
103 lines (83 loc) · 3.41 KB
/
train_test_functions.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
from utils import *
def train(epoch, model, train_loader, optimizer, args):
"""
:param epoch: current epoch
:param model: infoCatVAE model being trained
:param args: args associated to current infoCatVAE
:return:
"""
model.train()
train_loss, train_reco_loss, train_negH, train_KLD = 0, 0, 0, 0
class_loss = []
for batch_idx, (data,_) in enumerate(train_loader):
data = Variable(data)
if args.cuda:
data = data.cuda()
optimizer.zero_grad()
# Categorical VAE part
if args.is_convolutional:
recon_batch, mu, logvar, a, allmu, allvar = model(data.unsqueeze(1))
else:
recon_batch, mu, logvar, a, allmu, allvar = model(data)
# Information maximisation part
sample, labels = sampling(model, 10, args.mupriorT)
if args.cuda:
sample = sample.cuda()
labels = labels.cuda()
_, _, a_sample, _, _ = model.encode(sample)
# Loss computation + backpropagation
loss, la, lb, lc, mse_loss = loss_function(model, recon_batch.view(recon_batch.size(0), -1),
data.view(data.size(0), -1), a, allmu, allvar, args.mupriorT)
adv_loss = F.cross_entropy(a_sample, labels)
(loss + 100*adv_loss).backward()
train_loss += loss.data[0]
train_reco_loss += la.data[0]
train_negH += lb.data[0]
train_KLD += lc.data[0]
if args.cuda:
del sample
del labels
torch.cuda.empty_cache()
class_loss.append(adv_loss.data[0])
optimizer.step()
_, preds = torch.max(a, 1)
print('====> Epoch: {}, Reco: {:.2f}, negH: {:.2f}, KLD: {:.2f}, Class: {:.2f}'.format(
epoch,
train_reco_loss / len(train_loader.dataset),
train_negH / len(train_loader.dataset),
train_KLD / len(train_loader.dataset),
np.mean(class_loss)))
def test(epoch, model, test_loader, args):
"""
:param epoch: current epoch
:param model: infoCatVAE model being trained
:param args: args associated to current infoCatVAE
:return:
"""
global test_lost_list
model.eval()
test_loss = 0
for i, (data,_) in enumerate(test_loader):
data = Variable(data, volatile=True)
if args.cuda:
data = data.cuda()
if args.is_convolutional:
recon_batch, mu, logvar, a, allmu, allvar = model(data.unsqueeze(1))
else:
recon_batch, mu, logvar, a, allmu, allvar = model(data)
_, preds = torch.max(a, 1)
loss, la, lb, lc, mse_loss = loss_function(model, recon_batch.view(recon_batch.size(0), -1),
data.view(data.size(0), -1), a, allmu, allvar, args.mupriorT)
test_loss += loss.data[0]
if i == 0:
n = min(data.size(0), 10)
comparison = torch.cat([data.view(data.size(0), 1, 28, 28)[:n],
recon_batch.view(recon_batch.size(0), 1, 28, 28)[:n]])
save_image(comparison.data.cpu(),
'../InfoCatVAE' + str(epoch) + '.png', nrow=n)
if args.cuda:
del data
torch.cuda.empty_cache()
test_loss /= len(test_loader.dataset)
print('====> Test set loss: {:.4f}'.format(test_loss))
test_lost_list.append(test_loss)