| | |
|---|
| | | default=0.001, type=float) |
|---|
| | | parser.add_argument('--iter_ref', dest='iter_ref', help='Number of iterative refinement passes.', |
|---|
| | | default=1, type=int) |
|---|
| | | parser.add_argument('--dataset', dest='dataset', help='Dataset type.', default='Pose_300W_LP', type=str) |
|---|
| | | args = parser.parse_args() |
|---|
| | | return args |
|---|
| | | |
|---|
| | |
|---|
| | | transforms.RandomCrop(224), transforms.ToTensor(), |
|---|
| | | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) |
|---|
| | | |
|---|
| | | pose_dataset = datasets.Pose_300W_LP(args.data_dir, args.filename_list, |
|---|
| | | transformations) |
|---|
| | | if args.dataset == 'Pose_300W_LP': |
|---|
| | | pose_dataset = datasets.Pose_300W_LP(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': |
|---|
| | | pose_dataset = datasets.BIWI(args.data_dir, args.filename_list, transformations) |
|---|
| | | elif args.dataset == 'AFLW': |
|---|
| | | pose_dataset = datasets.AFLW(args.data_dir, args.filename_list, transformations) |
|---|
| | | elif args.dataset == 'AFW': |
|---|
| | | pose_dataset = datasets.AFW(args.data_dir, args.filename_list, transformations) |
|---|
| | | else: |
|---|
| | | print 'Error: not a valid dataset name' |
|---|
| | | sys.exit() |
|---|
| | | train_loader = torch.utils.data.DataLoader(dataset=pose_dataset, |
|---|
| | | batch_size=batch_size, |
|---|
| | | shuffle=True, |
|---|
| | |
|---|
| | | loss_pitch += alpha * loss_reg_pitch |
|---|
| | | loss_roll += alpha * loss_reg_roll |
|---|
| | | |
|---|
| | | loss_yaw *= 0.35 |
|---|
| | | |
|---|
| | | # Finetuning loss |
|---|
| | | loss_seq = [loss_yaw, loss_pitch, loss_roll] |
|---|
| | | for idx in xrange(args.iter_ref+1): |
|---|
| | | loss_angles = reg_criterion(angles[idx], label_angles.float()) |
|---|
| | | for idx in xrange(1,len(angles)): |
|---|
| | | label_angles_residuals = label_angles.float() - angles[0] |
|---|
| | | label_angles_residuals = label_angles_residuals.detach() |
|---|
| | | loss_angles = reg_criterion(angles[idx], label_angles_residuals) |
|---|
| | | loss_seq.append(loss_angles) |
|---|
| | | |
|---|
| | | grad_seq = [torch.Tensor(1).cuda(gpu) for _ in range(len(loss_seq))] |
|---|
