def train(train_queue, model, criterion, optimizer):
objs = utils.AvgRageCtrl()
face = utils.AvgRageCtrl()
rig_avg = utils.AvgRageCtrl()
reg = utils.AvgRageCtrl()
model.train()
for step, (input, label, rig, bbox) in enumerate(train_queue):
input = Variable(input, requires_grad=False).cuda()
label = Variable(label.float(), requires_grad=False).cuda(async=True)
bbox = Variable(bbox.float(), requires_grad=False).cuda(async=True)
rig = Variable(rig.float(), requires_grad=False).cuda(async=True)
optimizer.zero_grad()
label_pred, bbx_pred, rig_pred = model(input)
loss = criterion(label, label_pred, bbox, bbx_pred, rig, rig_pred)
loss.backward()
nn.utils.clip_grad_norm(model.parameters(), args.grad_clip)
optimizer.step()
prec1 = utils.compute_sigmoid_accuracy(label_pred, label)
prec2 = utils.compute_rig_accuracy(label, rig_pred, rig)
prec3 = criterion.box_loss(label, bbox, bbx_pred)
n = input.size(0)
objs.update(loss.data[0], n)
face.update(prec1.data[0], n)
rig_avg.update(prec2.data[0], n)
reg.update(prec3.data[0], n)
if step % args.report_freq == 0:
logger.info('train %03d %e %f %f %f', step, objs.avg, face.avg,
rig_avg.avg, reg.avg)
return face.avg, rig_avg.avg, reg.avg, objs.avg
def infer(valid_queue, model, criterion):
objs = utils.AvgRageCtrl()
face = utils.AvgRageCtrl()
rig_avg = utils.AvgRageCtrl()
reg = utils.AvgRageCtrl()
model.eval()
for step, (input, label, rig, bbox) in enumerate(valid_queue):
input = Variable(input, volatile=True).cuda()
label = Variable(label.float(), volatile=True).cuda(async=True)
bbox = Variable(bbox.float(), volatile=True).cuda(async=True)
rig = Variable(rig.float(), volatile=True).cuda(async=True)
label_pred, bbx_pred, rig_pred = model(input)
loss = criterion(label, label_pred, bbox, bbx_pred, rig, rig_pred)
prec1 = utils.compute_sigmoid_accuracy(label_pred, label)
prec2 = utils.compute_rig_accuracy_3(label, rig_pred, rig)
prec3 = criterion.box_loss(label, bbox, bbx_pred)
n = input.size(0)
objs.update(loss.data[0], n)
face.update(prec1.data[0], n)
rig_avg.update(prec2.data[0], n)
reg.update(prec3.data[0], n)
if step % args.report_freq == 0:
logging.info('valid %03d %e %f %f %f', step,
objs.avg, face.avg, rig_avg.avg, reg.avg)
return face.avg, rig_avg.avg, reg.avg, objs.avg
def train(train_queue, search_queue, model, architect, criterion, optimizer,
lr):
objs = utils.AvgRageCtrl()
face = utils.AvgRageCtrl()
rig_avg = utils.AvgRageCtrl()
reg = utils.AvgRageCtrl()
grad = utils.AvgRageCtrl()
for step, (input, label, rig, bbox) in enumerate(train_queue):
model.train()
n = input.size(0)
input = Variable(input, requires_grad=False).cuda()
label = Variable(label.float(), requires_grad=False).cuda(async=True)
bbox = Variable(bbox.float(), requires_grad=False).cuda(async=True)
rig = Variable(rig.float(), requires_grad=False).cuda(async=True)
target = (label, bbox, rig)
input_search, label_search, rig_search, bbox_search = next(
iter(search_queue))
input_search = Variable(input_search, requires_grad=False).cuda()
label_search = Variable(label_search.float(),
requires_grad=False).cuda(async=True)
rig_search = Variable(rig_search.float(),
requires_grad=False).cuda(async=True)
bbox_search = Variable(bbox_search.float(),
requires_grad=False).cuda(async=True)
target_search = (label_search, bbox_search, rig_search)
arch_grad_norm = architect.step(input,
target,
input_search,
target_search,
lr,
optimizer,
unrolled=args.unrolled)
grad.update(arch_grad_norm)
optimizer.zero_grad()
label_pred, bbx_pred, rig_pred = model(input)
loss = criterion(label, label_pred, bbox, bbx_pred, rig, rig_pred)
loss.backward()
nn.utils.clip_grad_norm(model.parameters(), args.grad_clip)
optimizer.step()
prec1 = utils.compute_sigmoid_accuracy(label_pred, label)
if args.stage == 1:
prec2 = utils.compute_rig_accuracy(label, rig_pred, rig)
elif args.stage == 2:
prec2 = utils.compute_rig_accuracy_2(label, rig_pred, rig)
prec3 = criterion.box_loss(label, bbox, bbx_pred)
objs.update(loss.data[0], n)
face.update(prec1.data[0], n)
rig_avg.update(prec2.data[0], n)
reg.update(prec3.data[0], n)
if step % args.report_freq == 0:
logging.info('train %03d %e %f %f %f', step, objs.avg, face.avg,
rig_avg.avg, reg.avg)
return face.avg, rig_avg.avg, reg.avg, objs.avg, grad.avg