natanielruiz
2017-09-23 d9781608243c42d55ced0c1d529f25eb30aeceb9
code/train_finetune.py
@@ -165,12 +165,10 @@
    print 'Second phase of training (finetuning layer).'
    for epoch in range(num_epochs_ft):
        for i, (images, labels, name) in enumerate(train_loader):
        for i, (images, labels, cont_labels, name) in enumerate(train_loader):
            images = Variable(images.cuda(gpu))
            label_yaw = Variable(labels[:,0].cuda(gpu))
            label_pitch = Variable(labels[:,1].cuda(gpu))
            label_roll = Variable(labels[:,2].cuda(gpu))
            label_angles = Variable(labels[:,:3].cuda(gpu))
            label_angles = Variable(cont_labels[:,:3].cuda(gpu))
            optimizer.zero_grad()
            model.zero_grad()
@@ -180,7 +178,7 @@
            # Finetuning loss
            loss_seq = []
            for idx in xrange(1,len(angles)):
                label_angles_residuals = label_angles.float() - angles[0]
                label_angles_residuals = label_angles - (angles[0] * 3 - 99)
                label_angles_residuals = label_angles_residuals.detach()
                loss_angles = reg_criterion(angles[idx], label_angles_residuals)
                loss_seq.append(loss_angles)