def child_train(train_queue, model, optimizer, global_step, arch_pool, arch_pool_prob, criterion, log_interval=100):
objs = utils.AverageMeter()
top1 = utils.AverageMeter()
top5 = utils.AverageMeter()
model.train()
for step, (input, target) in enumerate(train_queue):
input = utils.move_to_cuda(input)
target = utils.move_to_cuda(target)
optimizer.zero_grad()
# sample an arch to train
arch = utils.sample_arch(arch_pool, arch_pool_prob)
logits = model(input, arch)
loss = criterion(logits, target)
loss.backward()
optimizer.step()
prec1, prec5 = utils.accuracy(logits, target, topk=(1, 5))
n = input.size(0)
objs.update(loss.data, n)
top1.update(prec1.data, n)
top5.update(prec5.data, n)
global_step += 1
if global_step % log_interval == 0:
logging.info('Train %03d loss %e top1 %f top5 %f', global_step, objs.avg, top1.avg, top5.avg)
logging.info('Arch: %s', ' '.join(map(str, arch)))
if global_step >= args.max_num_updates:
break
return top1.avg, objs.avg, global_step
def child_train(train_queue, model, optimizer, global_step, arch_pool,
arch_pool_prob, criterion):
objs = utils.AvgrageMeter()
top1 = utils.AvgrageMeter()
top5 = utils.AvgrageMeter()
model.train()
for step, (input, target) in enumerate(train_queue):
input = input.cuda().requires_grad_()
target = target.cuda()
optimizer.zero_grad()
# sample an arch to train
arch = utils.sample_arch(arch_pool, arch_pool_prob)
logits, aux_logits = model(input, arch, global_step)
global_step += 1
loss = criterion(logits, target)
if aux_logits is not None:
aux_loss = criterion(aux_logits, target)
loss += 0.4 * aux_loss
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), args.child_grad_bound)
optimizer.step()
prec1, prec5 = utils.accuracy(logits, target, topk=(1, 5))
n = input.size(0)
objs.update(loss.data, n)
top1.update(prec1.data, n)
top5.update(prec5.data, n)
if (step + 1) % 100 == 0:
logging.info('Train %03d loss %e top1 %f top5 %f', step + 1,
objs.avg, top1.avg, top5.avg)
logging.info('Arch: %s', ' '.join(map(str, arch[0] + arch[1])))
return top1.avg, objs.avg, global_step