def train(model,
teacher_model,
optimizer,
loss_fn,
dataloader,
use_gpu,
epoch,
writer,
temp=1.0,
gamma=1.0):
model.train()
teacher_model.eval()
losses = utils.AverageMeter()
accuracies = utils.AverageMeter()
epoch_steps = len(dataloader)
teacher_model.eval()
for i, (train_batch, label_batch) in enumerate(dataloader):
niter = (epoch - 1) * epoch_steps + i
if use_gpu:
train_batch, label_batch = train_batch.cuda(), label_batch.cuda()
train_batch, label_batch = map(Variable, (train_batch, label_batch))
output_batch = model(train_batch)
teacher_output_batch = teacher_model(train_batch)
teacher_output_batch = Variable(teacher_output_batch)
loss = loss_fn(output_batch, label_batch, teacher_output_batch, temp,
gamma)
acc, _, _, _ = utils.find_metrics(output_batch, label_batch, use_gpu)
losses.update(loss.item())
accuracies.update(acc)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(("Step: {}, Current Loss: {}, RunningLoss: {}").format(
i, loss, losses.avg))
writer.add_scalar('data/stepwise_training_loss', losses.val, niter)
writer.add_scalar('data/stepwise_training_accuracy', accuracies.val,
niter)
writer.add_scalar('data/training_loss', losses.avg, epoch)
writer.add_scalar('data/training_accuracy', accuracies.avg, epoch)
return losses.avg
def validate(model, loss_fn, acc_fn, dataloader, use_gpu, epoch, writer):
losses = utils.AverageMeter()
accuracies = utils.AverageMeter()
precisions = utils.AverageMeter()
recalls = utils.AverageMeter()
model.eval()
for i, (train_batch, label_batch) in enumerate(dataloader):
if use_gpu:
train_batch, label_batch = train_batch.cuda(), label_batch.cuda()
with torch.no_grad():
train_batch, label_batch = Variable(
train_batch), Variable(label_batch)
output_batch = model(train_batch)
loss = loss_fn(output_batch, label_batch)
acc, prec, rec, _ = utils.find_metrics(output_batch, label_batch, use_gpu)
losses.update(loss.item())
accuracies.update(acc)
precisions.update(prec)
recalls.update(rec)
print(("Step: {}, Current Loss: {}, RunningLoss: {}").format(
i, loss, losses.avg))
writer.add_scalar('data/val_loss', losses.avg, epoch)
writer.add_scalar('data/val_accuracy', accuracies.avg, epoch)
writer.add_scalar('data/val_precision', precisions.avg, epoch)
writer.add_scalar('data/val_recall', recalls.avg, epoch)
return losses.avg
def train(model,
optimizer,
loss_fn,
acc_fn,
dataloader,
use_gpu,
epoch,
writer,
mixup=False,
alpha=1.0,
cutmix=False,
cutmix_prob=0.5,
beta=1.0):
model.train()
losses = utils.AverageMeter()
accuracies = utils.AverageMeter()
epoch_steps = len(dataloader)
for i, (train_batch, label_batch) in enumerate(dataloader):
niter = (epoch - 1) * epoch_steps + i
if use_gpu:
train_batch, label_batch = train_batch.cuda(), label_batch.cuda()
if mixup:
train_batch, label_batch_a, label_batch_b, lam = utils.mixup_data(
train_batch, label_batch, alpha, use_gpu)
train_batch, label_batch_a, label_batch_b = map(
Variable, (train_batch, label_batch_a, label_batch_b))
output_batch = model(train_batch)
loss = utils.mixed_loss_fn(loss_fn, output_batch, label_batch_a,
label_batch_b, lam)
losses.update(loss.item())
writer.add_scalar('data/stepwise_lambda', lam, niter)
elif cutmix:
train_batch, label_batch_a, label_batch_b, lam = utils.cutmix_data(
train_batch, label_batch, beta, cutmix_prob, use_gpu)
train_batch, label_batch_a, label_batch_b = map(
Variable, (train_batch, label_batch_a, label_batch_b))
output_batch = model(train_batch)
loss = utils.mixed_loss_fn(loss_fn, output_batch, label_batch_a,
label_batch_b, lam)
losses.update(loss.item())
writer.add_scalar('data/stepwise_lambda', lam, niter)
else:
# To set grad to zero on these
train_batch, label_batch = map(Variable,
(train_batch, label_batch))
output_batch = model(train_batch)
loss = loss_fn(output_batch, label_batch)
acc, _, _, _ = utils.find_metrics(output_batch, label_batch,
use_gpu)
losses.update(loss.item())
accuracies.update(acc)
writer.add_scalar('data/stepwise_training_accuracy',
accuracies.val, niter)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(("Step: {}, Current Loss: {}, RunningLoss: {}").format(
i, loss, losses.avg))
writer.add_scalar('data/stepwise_training_loss', losses.val, niter)
writer.add_scalar('data/training_loss', losses.avg, epoch)
if not (mixup or cutmix):
writer.add_scalar('data/training_accuracy', accuracies.avg, epoch)
return losses.avg