def evaluate_iwbo(model, eval_loader):
model = model.eval()
with torch.no_grad():
loss_sum = 0.0
count = 0
for i, x in enumerate(eval_loader):
x = x.to(args.device)
loss_sum += -iwbo_nats(model, x, k=args.iwbo_k, kbs=args.iwbo_kbs
).cpu().item() * len(x) / 50
count += len(x)
print('IWBO({}), Iter: {}/{}, PPLL: {:.3f}'.format(
args.iwbo_k, i + 1, len(eval_loader), loss_sum / count),
end='\r')
return loss_sum / count
loss = -model.log_prob(x.to(device)).mean()
loss.backward()
optimizer.step()
l += loss.detach().cpu().item()
print('Epoch: {}/{}, Iter: {}/{}, Nats: {:.3f}'.format(epoch+1, epochs, i+1, len(train_loader), l/(i+1)), end='\r')
print('')
##########
## Test ##
##########
print('Testing...')
with torch.no_grad():
l = 0.0
for i, x in enumerate(test_loader):
loss = iwbo_nats(model, x.to(device), k=10)
l += loss.detach().cpu().item()
print('Iter: {}/{}, Nats: {:.3f}'.format(i+1, len(test_loader), l/(i+1)), end='\r')
print('')
############
## Sample ##
############
print('Sampling...')
img = next(iter(test_loader))[:64]
samples = model.sample(64)
vutils.save_image(img.cpu().float(), './examples/results/mnist_data.png', nrow=8)
vutils.save_image(samples.cpu().float(), './examples/results/mnist_ndp_392.png', nrow=8)
print('Epoch: {}/{}, Iter: {}/{}, Nats: {:.3f}'.format(
epoch + 1, args.epochs, i + 1, len(train_loader),
loss_sum / (i + 1)),
end='\r')
print('')
final_train_bpd = loss_sum / len(train_loader)
#############
## Testing ##
#############
print('Testing...')
with torch.no_grad():
loss_sum = 0.0
for i, x in enumerate(test_loader):
loss = iwbo_nats(model, x.to(args.device), k=args.k, kbs=args.kbs)
loss_sum += loss.detach().cpu().item()
print('Iter: {}/{}, Nats: {:.3f}'.format(i + 1, len(test_loader),
loss_sum / (i + 1)),
end='\r')
print('')
final_test_nats = loss_sum / len(test_loader)
##############
## Sampling ##
##############
print('Sampling...')
if args.dataset == 'face_einstein':
bounds = [[0, 1], [0, 1]]
else: