From f111cb002b9c6065fdf6bb274ce5857a9e875e8c Mon Sep 17 00:00:00 2001 From: chenshijun <csj_sky@126.com> Date: 星期三, 05 六月 2019 15:38:49 +0800 Subject: [PATCH] face rectangle --- code/train_hopenet.py | 10 +++++----- 1 files changed, 5 insertions(+), 5 deletions(-) diff --git a/code/train_hopenet.py b/code/train_hopenet.py index 6278a87..56b9ca3 100644 --- a/code/train_hopenet.py +++ b/code/train_hopenet.py @@ -96,7 +96,7 @@ saved_state_dict = torch.load(args.snapshot) model.load_state_dict(saved_state_dict) - print 'Loading data.' + print('Loading data.') transformations = transforms.Compose([transforms.Scale(240), transforms.RandomCrop(224), transforms.ToTensor(), @@ -119,7 +119,7 @@ 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, @@ -142,7 +142,7 @@ {'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) @@ -184,7 +184,7 @@ 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))] + grad_seq = [torch.ones(1).cuda(gpu) for _ in range(len(loss_seq))] optimizer.zero_grad() torch.autograd.backward(loss_seq, grad_seq) optimizer.step() @@ -195,6 +195,6 @@ # 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') -- Gitblit v1.8.0