| | |
| | | # ResNet50 structure |
| | | model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], 66) |
| | | |
| | | print 'Loading snapshot.' |
| | | print('Loading snapshot.') |
| | | # Load snapshot |
| | | saved_state_dict = torch.load(snapshot_path) |
| | | model.load_state_dict(saved_state_dict) |
| | | |
| | | print 'Loading data.' |
| | | print('Loading data.') |
| | | |
| | | transformations = transforms.Compose([transforms.Scale(224), |
| | | transforms.CenterCrop(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() |
| | | test_loader = torch.utils.data.DataLoader(dataset=pose_dataset, |
| | | batch_size=args.batch_size, |
| | |
| | | |
| | | model.cuda(gpu) |
| | | |
| | | print 'Ready to test network.' |
| | | print('Ready to test network.') |
| | | |
| | | # Test the Model |
| | | model.eval() # Change model to 'eval' mode (BN uses moving mean/var). |