| | |
|---|
| | | model = hopenet.AlexNet(66) |
|---|
| | | load_filtered_state_dict(model, model_zoo.load_url(model_urls['alexnet'])) |
|---|
| | | |
|---|
| | | print 'Loading data.' |
|---|
| | | print('Loading data.') |
|---|
| | | |
|---|
| | | transformations = transforms.Compose([transforms.Scale(240), |
|---|
| | | transforms.RandomCrop(224), transforms.ToTensor(), |
|---|
| | |
|---|
| | | elif args.dataset == 'AFW': |
|---|
| | | pose_dataset = datasets.AFW(args.data_dir, args.filename_list, transformations) |
|---|
| | | else: |
|---|
| | | print 'Error: not a valid dataset name' |
|---|
| | | print('Error: not a valid dataset name') |
|---|
| | | sys.exit() |
|---|
| | | train_loader = torch.utils.data.DataLoader(dataset=pose_dataset, |
|---|
| | | batch_size=batch_size, |
|---|
| | |
|---|
| | | # Regression loss coefficient |
|---|
| | | alpha = args.alpha |
|---|
| | | |
|---|
| | | idx_tensor = [idx for idx in xrange(66)] |
|---|
| | | idx_tensor = Variable(torch.FloatTensor(idx_tensor)).cuda(gpu) |
|---|
| | | |
|---|
| | | optimizer = torch.optim.Adam([{'params': get_ignored_params(model), 'lr': 0}, |
|---|
| | | {'params': get_non_ignored_params(model), 'lr': args.lr}, |
|---|
| | | {'params': get_fc_params(model), 'lr': args.lr * 5}], |
|---|
| | | lr = args.lr) |
|---|
| | | |
|---|
| | | print 'Ready to train network.' |
|---|
| | | print('Ready to train network.') |
|---|
| | | for epoch in range(num_epochs): |
|---|
| | | for i, (images, labels, cont_labels, name) in enumerate(train_loader): |
|---|
| | | images = Variable(images).cuda(gpu) |
|---|
| | |
|---|
| | | label_roll_cont = Variable(cont_labels[:,2]).cuda(gpu) |
|---|
| | | |
|---|
| | | # Forward pass |
|---|
| | | yaw, pitch, roll, angles = model(images) |
|---|
| | | pre_yaw, pre_pitch, pre_roll = model(images) |
|---|
| | | |
|---|
| | | # Cross entropy loss |
|---|
| | | loss_yaw = criterion(yaw, label_yaw) |
|---|
| | | loss_pitch = criterion(pitch, label_pitch) |
|---|
| | | loss_roll = criterion(roll, label_roll) |
|---|
| | | loss_yaw = criterion(pre_yaw, label_yaw) |
|---|
| | | loss_pitch = criterion(pre_pitch, label_pitch) |
|---|
| | | loss_roll = criterion(pre_roll, label_roll) |
|---|
| | | |
|---|
| | | # MSE loss |
|---|
| | | yaw_predicted = angles[:,0] |
|---|
| | | pitch_predicted = angles[:,1] |
|---|
| | | roll_predicted = angles[:,2] |
|---|
| | | yaw_predicted = softmax(pre_yaw) |
|---|
| | | pitch_predicted = softmax(pre_pitch) |
|---|
| | | roll_predicted = softmax(pre_roll) |
|---|
| | | |
|---|
| | | yaw_predicted = torch.sum(yaw_predicted * idx_tensor, 1) * 3 - 99 |
|---|
| | | pitch_predicted = torch.sum(pitch_predicted * idx_tensor, 1) * 3 - 99 |
|---|
| | | roll_predicted = torch.sum(roll_predicted * idx_tensor, 1) * 3 - 99 |
|---|
| | | |
|---|
| | | loss_reg_yaw = reg_criterion(yaw_predicted, label_yaw_cont) |
|---|
| | | loss_reg_pitch = reg_criterion(pitch_predicted, label_pitch_cont) |
|---|
| | |
|---|
| | | loss_roll += alpha * loss_reg_roll |
|---|
| | | |
|---|
| | | loss_seq = [loss_yaw, loss_pitch, loss_roll] |
|---|
| | | grad_seq = [torch.Tensor(1).cuda(gpu) for _ in range(len(loss_seq))] |
|---|
| | | optimizer.zero_grad() |
|---|
| | | grad_seq = [torch.ones(1).cuda(gpu) for _ in range(len(loss_seq))] |
|---|
| | | torch.autograd.backward(loss_seq, grad_seq) |
|---|
| | | optimizer.step() |
|---|
| | | |
|---|
| | | if (i+1) % 100 == 0: |
|---|
| | | print ('Epoch [%d/%d], Iter [%d/%d] Losses: Yaw %.4f, Pitch %.4f, Roll %.4f' |
|---|
| | | print('Epoch [%d/%d], Iter [%d/%d] Losses: Yaw %.4f, Pitch %.4f, Roll %.4f' |
|---|
| | | %(epoch+1, num_epochs, i+1, len(pose_dataset)//batch_size, loss_yaw.data[0], loss_pitch.data[0], loss_roll.data[0])) |
|---|
| | | |
|---|
| | | # Save models at numbered epochs. |
|---|
| | | if epoch % 1 == 0 and epoch < num_epochs: |
|---|
| | | print 'Taking snapshot...' |
|---|
| | | print('Taking snapshot...') |
|---|
| | | torch.save(model.state_dict(), |
|---|
| | | 'output/snapshots/' + args.output_string + '_epoch_'+ str(epoch+1) + '.pkl') |
|---|
