示例#1
0
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
示例#2
0
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
示例#3
0
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