best_loss = checkpoint['train_loss']
print('\n******************** Best loss resumed is :', best_loss,
' ************************')
start_epoch = checkpoint['epoch'] + 1
# else:
# print(' # Loading pretrained model # ')
# posenet.load_state_dict(torch.load(opt.pretrained_model))
criterion = MultiTaskLoss(opt, config)
if use_cuda:
posenet = torch.nn.parallel.DataParallel(
posenet.cuda()) # , device_ids=[0, 1, 2, 3]
# module.cuda() only move the registered parameters to GPU.
criterion.cuda()
# 如果我们每次训练的输入数据的size不变，那么开启这个就会加快我们的训练速度
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=10,
gamma=0.2,
last_epoch=-1) # 设置学习率下降策略
for i in range(start_epoch):
# update the learning rate for start_epoch times
scheduler.step()
else:
raise NotImplementedError
# loader (batch_size MUST be 1)
loader = DataLoader(data_set,
batch_size=1,
shuffle=False,
num_workers=5,
pin_memory=True)
# activate GPUs
CUDA = torch.cuda.is_available()
torch.manual_seed(seed)
if CUDA:
torch.cuda.manual_seed(seed)
model.cuda()
# opencv init
fourcc = cv2.VideoWriter_fourcc(*'XVID')
model_name = 'posenet' if args.weights.find('posenet') >= 0 else 'vidvo'
out_filename = osp.join(
args.output_dir,
'{:s}_{:s}_attention_{:s}.avi'.format(args.dataset, args.scene,
model_name))
# get frame size
img, _ = data_set[0]
vwrite = cv2.VideoWriter(out_filename,
fourcc=fourcc,
fps=20.0,
frameSize=(img.size(2), img.size(1)))
print 'Initialized VideoWriter to {:s} with frames size {:d} x {:d}'.\
state[k] = v.cuda()
print('\nOptimizer has been resumed from checkpoint...')
best_loss = checkpoint['train_loss']
print('\n******************** Best loss resumed is :', best_loss, ' ************************')
start_epoch = checkpoint['epoch'] + 1
# else:
# print(' # Loading pretrained model # ')
# posenet.load_state_dict(torch.load(opt.pretrained_model))
criterion = MultiTaskLoss(opt, config)
if use_cuda:
posenet = torch.nn.parallel.DataParallel(posenet.cuda()) # , device_ids=[0, 1, 2, 3]
# module.cuda() only move the registered parameters to GPU.
criterion.cuda()
torch.backends.cudnn.benchmark = True # 如果我们每次训练的输入数据的size不变，那么开启这个就会加快我们的训练速度
# torch.backends.cudnn.deterministic = True
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.2, last_epoch=-1) # 设置学习率下降策略
for i in range(start_epoch):
# update the learning rate for start_epoch times
scheduler.step()
for param in posenet.parameters():
if param.requires_grad:
print('Parameters of network: Autograd')