""" Run an epoch over the given loader """

if train:

net.train()

tracker_class, tracker_params = tracker.MovingMeanMonitor, {'momentum': 0.99}

else:

net.eval()

tracker_class, tracker_params = tracker.MeanMonitor, {}

answ = []

idxs = []

accs = []

pbar = tqdm(loader, desc=f'{prefix} E{epoch}')

loss_tracker = tracker.track('{}_loss'.format(prefix), tracker_class(**tracker_params))

acc_tracker = tracker.track('{}_acc'.format(prefix), tracker_class(**tracker_params))

log_softmax = nn.LogSoftmax().cuda()

for v, q, a, idx, q_len in pbar:

var_params = {

'volatile': not train,

'requires_grad': False,

}

v = Variable(v.cuda(), **var_params)

q = Variable(q.cuda(), **var_params)

a = Variable(a.cuda(), **var_params)

q_len = Variable(q_len.cuda(), **var_params)

out = net(v, q, q_len)

nll = -log_softmax(out, dim=1)

loss = (nll * a / 10).sum(dim=1).mean()

acc = utils.batch_accuracy(out.data, a.data).cpu()

if train:

global total_iterations

update_learning_rate(optimizer, total_iterations)

optimizer.zero_grad()

loss.backward()

optimizer.step()

total_iterations = 1

else:

# store information about evaluation of this minibatch

_, answer = out.data.cpu().max(dim=1)

answ.append(answer.view(-1))

accs.append(acc.view(-1))

idxs.append(idx.view(-1).clone())

loss_tracker.append(loss.item())

# acc_tracker.append(acc.mean())

for a in acc:

acc_tracker.append(a.item())

fmt = '{:.4f}'.format

pbar.set_postfix(loss=fmt(loss_tracker.mean.value), acc=fmt(acc_tracker.mean.value))

if not train:

answ = list(torch.cat(answ, dim=0))

accs = list(torch.cat(accs, dim=0))

idxs = list(torch.cat(idxs, dim=0))

if len(sys.argv) > 1:

name = ' '.join(sys.argv[1:])

else:

from datetime import datetime

name = datetime.now().strftime("%Y-%m-%d_%H:%M:%S")

target_name = os.path.join('logs', '{}.pth'.format(name))

print('will save to {}'.format(target_name))

cudnn.benchmark = True

train_loader = data.get_loader(train=True)

val_loader = data.get_loader(val=True)

net = nn.DataParallel(model.Net(train_loader.dataset.num_tokens)).cuda()

optimizer = optim.Adam([p for p in net.parameters() if p.requires_grad])

tracker = utils.Tracker()

config_as_dict = {k: v for k, v in vars(config).items() if not k.startswith('__')}

for i in range(config.epochs):

_ = run(net, train_loader, optimizer, tracker, train=True, prefix='train', epoch=i)

r = run(net, val_loader, optimizer, tracker, train=False, prefix='val', epoch=i)

results = {

'name': name,

'tracker': tracker.to_dict(),

'config': config_as_dict,

'weights': net.state_dict(),

'eval': {

'answers': r[0],

'accuracies': r[1],

'idx': r[2],

},

'vocab': train_loader.dataset.vocab,

}