| | |
|---|
| | | model = hopenet.ResNet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], 3) |
|---|
| | | load_filtered_state_dict(model, model_zoo.load_url('https://download.pytorch.org/models/resnet50-19c8e357.pth')) |
|---|
| | | |
|---|
| | | 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, |
|---|
| | |
|---|
| | | {'params': get_fc_params(model), 'lr': args.lr * 5}], |
|---|
| | | lr = args.lr) |
|---|
| | | |
|---|
| | | print 'Ready to train network.' |
|---|
| | | print 'First phase of training.' |
|---|
| | | print('Ready to train network.') |
|---|
| | | print('First phase of training.') |
|---|
| | | for epoch in range(num_epochs): |
|---|
| | | for i, (images, labels, cont_labels, name) in enumerate(train_loader): |
|---|
| | | images = Variable(images).cuda(gpu) |
|---|
| | |
|---|
| | | |
|---|
| | | # 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') |
|---|
