natanielruiz
2017-09-27 43416c4717d2430c3e11f042294d12b781fee2e1
code/train_preangles.py
@@ -125,6 +125,8 @@
    if args.dataset == 'Pose_300W_LP':
        pose_dataset = datasets.Pose_300W_LP(args.data_dir, args.filename_list, transformations)
    if args.dataset == 'Pose_300W_LP_random_ds':
        pose_dataset = datasets.Pose_300W_LP_random_ds(args.data_dir, args.filename_list, transformations)
    elif args.dataset == 'AFLW2000':
        pose_dataset = datasets.AFLW2000(args.data_dir, args.filename_list, transformations)
    elif args.dataset == 'BIWI':
@@ -161,10 +163,10 @@
    print 'Ready to train network.'
    print 'First phase of training.'
    for epoch in range(num_epochs):
        start = time.time()
        # start = time.time()
        for i, (images, labels, cont_labels, name) in enumerate(train_loader):
            print i
            print 'start: ', time.time() - start
            # print i
            # print 'start: ', time.time() - start
            images = Variable(images).cuda(gpu)
            label_yaw = Variable(labels[:,0]).cuda(gpu)
            label_pitch = Variable(labels[:,1]).cuda(gpu)
@@ -207,7 +209,7 @@
            torch.autograd.backward(loss_seq, grad_seq)
            optimizer.step()
            print 'end: ', time.time() - start
            # print 'end: ', time.time() - start
            if (i+1) % 100 == 0:
                print ('Epoch [%d/%d], Iter [%d/%d] Losses: Yaw %.4f, Pitch %.4f, Roll %.4f'