def main():
global constants
# enable cudnn auto-tuner to find the best algorithm for the given harware
cudnn.benchmark = True
sm = SessionManager(dataset_name=args.dataset_name,
resume_id=args.session_id)
labeled_trainloader, unlabeled_trainloader, val_loader, test_loader, class_names, constants = \
sm.load_dataset(args)
ts, writer = sm.load_checkpoint(class_names, constants)
step = 0
# Train and val
for epoch in range(ts.start_epoch, constants['epochs']):
print('\nEpoch: [%d | %d]' % (epoch + 1, constants['epochs']))
step = constants['train_iteration'] * (epoch + 1)
if constants['enable_mixmatch']:
train_loss = train(labeled_trainloader, unlabeled_trainloader,
epoch, ts)
else:
train_loss = train_supervised(labeled_trainloader, epoch, ts)
losses, accs, confs, names = validate_all(labeled_trainloader,
val_loader, test_loader,
train_loss, ts)
tensorboard_write(writer, losses, accs, confs, names, class_names,
step)
# save model and other training variables
sm.save_checkpoint(accs[names['validation']], epoch)
sm.close()
