def train(epoch):
print('\nEpoch: %d' % epoch)
model.train()
train_loss, correct, total = 0, 0, 0
for batch_idx, (inputs, targets) in enumerate(train_loader):
inputs, targets = inputs.to('cuda'), targets.to('cuda')
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
train_loss += loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
progress_bar(
batch_idx, len(train_loader),
'Loss: %.3f | Acc: %.3f%% (%d/%d)' %
(train_loss /
(batch_idx + 1), 100. * correct / total, correct, total))
if batch_idx % cfg.log_interval == 0: #every log_interval mini_batches...
summary_writer.add_scalar(
'Loss/train', train_loss / (batch_idx + 1),
epoch * len(train_loader) + batch_idx)
summary_writer.add_scalar(
'Accuracy/train', 100. * correct / total,
epoch * len(train_loader) + batch_idx)
summary_writer.add_scalar(
'learning rate', optimizer.param_groups[0]['lr'],
epoch * len(train_loader) + batch_idx)
def test(epoch):
global best_acc
model.eval()
test_loss, correct, total = 0, 0, 0
with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(eval_loader):
inputs, targets = inputs.to(device), targets.to(device)
outputs = model(inputs)
loss = criterion(outputs, targets)
test_loss += loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
progress_bar(batch_idx, len(eval_loader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
% (test_loss / (batch_idx + 1), 100. * correct / total, correct, total))
if batch_idx % cfg.log_interval == 0: # every log_interval mini_batches...
summary_writer.add_scalar('Loss/test', test_loss / (batch_idx + 1),
epoch * len(train_loader) + batch_idx)
summary_writer.add_scalar('Accuracy/test', 100. * correct / total,
epoch * len(train_loader) + batch_idx)
acc = 100. * correct / total
if acc > best_acc:
print('Saving Models..')
state = {
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'acc': acc,
'epoch': epoch,
}
torch.save(state, os.path.join(cfg.ckpt_dir, '{}_checkpoint.t7'.format(cfg.log_name)))
best_acc = acc