| | |
| | | |
| | | print 'Loading data.' |
| | | |
| | | transformations = transforms.Compose([transforms.Scale(224),transforms.RandomCrop(224), |
| | | transforms.ToTensor()]) |
| | | # transformations = transforms.Compose([transforms.Scale(224), |
| | | # transforms.RandomCrop(224), |
| | | # transforms.ToTensor()]) |
| | | |
| | | transformations = transforms.Compose([transforms.Scale(250), |
| | | 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) |
| | |
| | | optimizer = torch.optim.Adam([{'params': get_ignored_params(model), 'lr': args.lr}, |
| | | {'params': get_non_ignored_params(model), 'lr': args.lr * 10}], |
| | | lr = args.lr) |
| | | # optimizer = torch.optim.SGD([{'params': get_ignored_params(model), 'lr': args.lr}, |
| | | # {'params': get_non_ignored_params(model), 'lr': args.lr}], |
| | | # lr = args.lr, |
| | | # momentum = 0.9, weight_decay=0.01) |
| | | |
| | | print 'Ready to train network.' |
| | | |
| | |
| | | if (i+1) % 100 == 0: |
| | | 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])) |
| | | if epoch == 0: |
| | | torch.save(model.state_dict(), |
| | | 'output/snapshots/resnet50_iter_'+ str(i+1) + '.pkl') |
| | | # if epoch == 0: |
| | | # torch.save(model.state_dict(), |
| | | # 'output/snapshots/resnet50_lbatch_iter_'+ str(i+1) + '.pkl') |
| | | |
| | | # Save models at numbered epochs. |
| | | if epoch % 1 == 0 and epoch < num_epochs - 1: |
| | | print 'Taking snapshot...' |
| | | torch.save(model.state_dict(), |
| | | 'output/snapshots/resnet50_epoch_'+ str(epoch+1) + '.pkl') |
| | | 'output/snapshots/resnet50_norm_norot_epoch_'+ str(epoch+1) + '.pkl') |
| | | |
| | | # Save the final Trained Model |
| | | torch.save(model.state_dict(), 'output/snapshots/resnet50_epoch' + str(epoch+1) + '.pkl') |
| | | torch.save(model.state_dict(), 'output/snapshots/resnet50_norm_norot_epoch_' + str(epoch+1) + '.pkl') |