def main():
dataset_path = '/home/emannuell/Documentos/mestrado/dataset/data_new/'
output_path = 'output/convergeTest'
train_set = GazeDataset(root_dir=dataset_path, training='train')
train_data_loader = DataLoader(train_set,
batch_size=10,
shuffle=True,
num_workers=4)
test_set = GazeDataset(root_dir=dataset_path, training='test')
test_data_loader = DataLoader(test_set,
batch_size=2,
shuffle=False,
num_workers=4)
net = GazeNet()
net = DataParallel(net)
net.cuda()
resume_training = False
if resume_training:
pretrained_dict = torch.load('model/pretrained_model.pkl')
model_dict = net.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
net.load_state_dict(model_dict)
test(net, test_data_loader)
exit()
method = 'Adam'
# 0.0001
learning_rate = 0.001
optimizer_s1 = optim.Adam(
[{
'params': net.module.head_pose_transform.parameters(),
'initial_lr': learning_rate
}, {
'params': net.module.eye_position_transform.parameters(),
'initial_lr': learning_rate
}, {
'params': net.module.fusion.parameters(),
'initial_lr': learning_rate
}],
lr=learning_rate,
weight_decay=0.0001)
optimizer_s2 = optim.Adam([{
'params': net.module.fpn_net.parameters(),
'initial_lr': learning_rate
}],
lr=learning_rate,
weight_decay=0.0001)
optimizer_s3 = optim.Adam([{
'params': net.parameters(),
'initial_lr': learning_rate
}],
lr=learning_rate * 0.1,
weight_decay=0.0001)
lr_scheduler_s1 = optim.lr_scheduler.StepLR(optimizer_s1,
step_size=5,
gamma=0.1,
last_epoch=-1)
lr_scheduler_s2 = optim.lr_scheduler.StepLR(optimizer_s2,
step_size=5,
gamma=0.1,
last_epoch=-1)
lr_scheduler_s3 = optim.lr_scheduler.StepLR(optimizer_s3,
step_size=5,
gamma=0.1,
last_epoch=-1)
max_epoch = 25
epoch = 0
while epoch < max_epoch:
if epoch == 0:
lr_scheduler = lr_scheduler_s1
optimizer = optimizer_s1
# 5
elif epoch == 15:
lr_scheduler = lr_scheduler_s2
optimizer = optimizer_s2
# 9
elif epoch == 20:
lr_scheduler = lr_scheduler_s3
optimizer = optimizer_s3
# optimizer.step()
lr_scheduler.step()
running_loss = []
ep_heatmap_loss = []
ep_m_angle_loss = []
for i, data in tqdm(enumerate(train_data_loader)):
image, gaze_field, eye_position, gt_position, gt_heatmap, head_pose = \
data['image'], data['gaze_field'], data['eye_position'], data['gt_position'], data['gt_heatmap'], data['head_pose']
image, gaze_field, eye_position, gt_position, gt_heatmap, head_pose = \
map(lambda x: Variable(x.cuda()), [image, gaze_field, eye_position, gt_position, gt_heatmap, head_pose])
optimizer.zero_grad()
direction, predict_heatmap = net(
[image, gaze_field, eye_position, head_pose])
heatmap_loss, m_angle_loss = F_loss(direction, predict_heatmap,
eye_position, gt_position,
gt_heatmap)
ep_heatmap_loss.append(np.array(heatmap_loss.cpu().data))
ep_m_angle_loss.append(np.array(m_angle_loss.cpu().data))
if epoch == 0:
loss = m_angle_loss
elif epoch >= 15 and epoch <= 20:
loss = heatmap_loss
else:
loss = m_angle_loss + heatmap_loss
loss.backward()
optimizer.step()
running_loss.append(
[heatmap_loss.data, m_angle_loss.data, loss.data])
# if i % 10 == 9:
# logging.info('%s %s %s'%(str(np.mean(running_loss, axis=0)), method, str(lr_scheduler.get_lr())))
# running_loss = []
epoch += 1
print('==== Training loss ====')
logging.info('Epoch: %s' % epoch)
logging.info('heatmap loss: %s' %
str(np.mean(np.array(ep_heatmap_loss))))
logging.info('mean angle loss: %s' %
str(np.mean(np.array(ep_m_angle_loss))))
logging.info('file: %s' % output_path +
'/epoch_{}_loss_{}.pkl'.format(epoch, loss.data))
if not os.path.exists(output_path):
os.makedirs(output_path)
print('Saving model to output path: ',
output_path + '/epoch_{}_loss_{}.pkl'.format(epoch, loss.data))
torch.save(
net.state_dict(),
output_path + '/epoch_{}_loss_{}.pkl'.format(epoch, loss.data))
def main():
train_set = GazeDataset(root_dir='../GazeFollowData/',
mat_file='../GazeFollowData/train_annotations.mat',
training='train')
train_data_loader = DataLoader(train_set,
batch_size=32 * 4,
shuffle=True,
num_workers=16)
test_set = GazeDataset(root_dir='../GazeFollowData/',
mat_file='../GazeFollowData/test2_annotations.mat',
training='test')
test_data_loader = DataLoader(test_set,
batch_size=32 * 4,
shuffle=False,
num_workers=8)
net = GazeNet()
net = DataParallel(net)
net.cuda()
resume_training = False
if resume_training:
pretrained_dict = torch.load('../model/pretrained_model.pkl')
model_dict = net.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
net.load_state_dict(model_dict)
test(net, test_data_loader)
exit()
method = 'Adam'
learning_rate = 0.0001
optimizer_s1 = optim.Adam(
[{
'params': net.module.face_net.parameters(),
'initial_lr': learning_rate
}, {
'params': net.module.face_process.parameters(),
'initial_lr': learning_rate
}, {
'params': net.module.eye_position_transform.parameters(),
'initial_lr': learning_rate
}, {
'params': net.module.fusion.parameters(),
'initial_lr': learning_rate
}],
lr=learning_rate,
weight_decay=0.0001)
optimizer_s2 = optim.Adam([{
'params': net.module.fpn_net.parameters(),
'initial_lr': learning_rate
}],
lr=learning_rate,
weight_decay=0.0001)
optimizer_s3 = optim.Adam([{
'params': net.parameters(),
'initial_lr': learning_rate
}],
lr=learning_rate * 0.1,
weight_decay=0.0001)
lr_scheduler_s1 = optim.lr_scheduler.StepLR(optimizer_s1,
step_size=5,
gamma=0.1,
last_epoch=-1)
lr_scheduler_s2 = optim.lr_scheduler.StepLR(optimizer_s2,
step_size=5,
gamma=0.1,
last_epoch=-1)
lr_scheduler_s3 = optim.lr_scheduler.StepLR(optimizer_s3,
step_size=5,
gamma=0.1,
last_epoch=-1)
max_epoch = 25
epoch = 0
while epoch < max_epoch:
if epoch == 0:
lr_scheduler = lr_scheduler_s1
optimizer = optimizer_s1
elif epoch == 7:
lr_scheduler = lr_scheduler_s2
optimizer = optimizer_s2
elif epoch == 15:
lr_scheduler = lr_scheduler_s3
optimizer = optimizer_s3
lr_scheduler.step()
running_loss = []
for i, data in tqdm(enumerate(train_data_loader)):
image, face_image, gaze_field, eye_position, gt_position, gt_heatmap = \
data['image'], data['face_image'], data['gaze_field'], data['eye_position'], data['gt_position'], data['gt_heatmap']
image, face_image, gaze_field, eye_position, gt_position, gt_heatmap = \
map(lambda x: Variable(x.cuda()), [image, face_image, gaze_field, eye_position, gt_position, gt_heatmap])
#for var in [image, face_image, gaze_field, eye_position, gt_position]:
# print var.shape
optimizer.zero_grad()
direction, predict_heatmap = net(
[image, face_image, gaze_field, eye_position])
heatmap_loss, m_angle_loss = \
F_loss(direction, predict_heatmap, eye_position, gt_position, gt_heatmap)
if epoch == 0:
loss = m_angle_loss
elif epoch >= 7 and epoch <= 14:
loss = heatmap_loss
else:
loss = m_angle_loss + heatmap_loss
loss.backward()
optimizer.step()
running_loss.append(
[heatmap_loss.data[0], m_angle_loss.data[0], loss.data[0]])
if i % 10 == 9:
logging.info('%s %s %s' % (str(np.mean(running_loss, axis=0)),
method, str(lr_scheduler.get_lr())))
running_loss = []
epoch += 1
save_path = '../model/two_stage_fpn_concat_multi_scale_' + method
if not os.path.exists(save_path):
os.makedirs(save_path)
torch.save(net.state_dict(),
save_path + '/model_epoch{}.pkl'.format(epoch))
test(net, test_data_loader)
def main():
'''
train_set = GazeDataset(root_dir='../../data/',
mat_file='../../data/train_annotations.mat',
training='train')
train_data_loader = DataLoader(train_set, batch_size=48,
shuffle=True, num_workers=8)
test_set = GazeDataset(root_dir='../../test_data/',
mat_file='../../test_data/test2_annotations.mat',
training='test')
test_data_loader = DataLoader(test_set, batch_size=32,
shuffle=False, num_workers=8)
'''
dis_train_sets = dataset_wrapper(
root_dir='../../data/',
mat_file='../../data/train_annotations.mat',
training='train')
#dis_train_data_loader = DataLoader(dis_train_sets[0], batch_size=48,
# shuffle=True, num_workers=8)
dis_test_sets = dataset_wrapper(
root_dir='../../test_data/',
mat_file='../../test_data/test2_annotations.mat',
training='test')
#dis_test_data_loader = DataLoader(dis_test_sets[0], batch_size=32,
# shuffle=False, num_workers=8)
dis_train_data_loaders, dis_test_data_loaders = [], []
for i in range(16):
dis_train_data_loaders.append(
DataLoader(dis_train_sets[i],
batch_size=40,
shuffle=True,
num_workers=8))
dis_test_data_loaders.append(
DataLoader(dis_test_sets[i],
batch_size=16,
shuffle=False,
num_workers=1))
net = GazeNet()
net = DataParallel(net)
net.cuda()
#print(next(net.module.fpn_net.parameters()).is_cuda)
##print(next(net.module.fpn_net.parameters()).is_cuda)
area_count = 8
area_in_network = int(16 / area_count)
cur_area_idx = 0
fpn_weights_transferred = False
for i in range(area_count):
net.module.change_fpn(i)
if not next(net.module.fpn_net.parameters()).is_cuda:
net.module.fpn_net.cuda()
net.module.change_fpn(cur_area_idx)
##print(next(net.module.fpn_net.parameters()).is_cuda)
#exit(0)
resume_training = False
if resume_training:
pretrained_dict = torch.load('../model/pretrained_model.pkl')
model_dict = net.state_dict()
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
net.load_state_dict(model_dict)
test(net, test_data_loader)
exit()
method = 'Adam'
learning_rate = 0.0001
optimizer_s1 = optim.Adam(
[{
'params': net.module.face_net.parameters(),
'initial_lr': learning_rate
}, {
'params': net.module.face_process.parameters(),
'initial_lr': learning_rate
}, {
'params': net.module.eye_position_transform.parameters(),
'initial_lr': learning_rate
}, {
'params': net.module.fusion.parameters(),
'initial_lr': learning_rate
}],
lr=learning_rate,
weight_decay=0.0001)
#optimizer_s2 = optim.Adam([{'params': net.module.fpn_net.parameters(),
# 'initial_lr': learning_rate}],
# lr=learning_rate, weight_decay=0.0001)
optimizer_s2s, optimizer_s3s = [], []
for i in range(area_count):
net.module.change_fpn(i)
optimizer_s2 = optim.Adam(
[{
'params': net.module.fpn_nets[i].parameters(),
'initial_lr': learning_rate
}],
lr=learning_rate,
weight_decay=0.0001)
optimizer_s3 = optim.Adam([{
'params': net.parameters(),
'initial_lr': learning_rate
}],
lr=learning_rate * 0.1,
weight_decay=0.0001)
optimizer_s2s.append(optimizer_s2)
optimizer_s3s.append(optimizer_s3)
optimizer_s2 = optimizer_s2s[0]
optimizer_s3 = optimizer_s3s[0]
lr_scheduler_s1 = optim.lr_scheduler.StepLR(optimizer_s1,
step_size=5,
gamma=0.1,
last_epoch=-1)
#lr_scheduler_s2 = optim.lr_scheduler.StepLR(optimizer_s2, step_size=5, gamma=0.1, last_epoch=-1)
lr_scheduler_s2s, lr_scheduler_s3s = [], []
for i in range(area_count):
lr_scheduler_s2 = optim.lr_scheduler.StepLR(optimizer_s2s[i],
step_size=5,
gamma=0.1,
last_epoch=-1)
lr_scheduler_s3 = optim.lr_scheduler.StepLR(optimizer_s3s[i],
step_size=5,
gamma=0.1,
last_epoch=-1)
lr_scheduler_s2s.append(lr_scheduler_s2)
lr_scheduler_s3s.append(lr_scheduler_s3)
lr_scheduler_s2 = lr_scheduler_s2s[0]
lr_scheduler_s3 = lr_scheduler_s3s[0]
# Set the model to use the first FPN
net.module.change_fpn(cur_area_idx)
max_epoch = 30
epoch = 0
#epoch = 7
while epoch < max_epoch:
logging.info('\n--- Epoch: %s\n' % str(epoch))
if epoch == 0:
lr_scheduler = lr_scheduler_s1
optimizer = optimizer_s1
elif epoch == 7:
lr_scheduler = lr_scheduler_s2
optimizer = optimizer_s2
elif epoch == 15:
lr_scheduler = lr_scheduler_s3
optimizer = optimizer_s3
#lr_scheduler.step()
#lr_scheduler.step()
running_loss = []
#for data_loader_idx in range(len(dis_train_data_loaders)):
for data_loader_idx in range(len(dis_train_data_loaders)):
train_data_loader = dis_train_data_loaders[data_loader_idx]
if epoch >= 10:
#if epoch >= 7:
if not fpn_weights_transferred:
net.module.transfer_fpn_weights()
fpn_weights_transferred = True
area_idx = int(data_loader_idx / area_in_network)
if cur_area_idx != area_idx:
cur_area_idx = area_idx
net.module.change_fpn(cur_area_idx)
if epoch < 15:
lr_scheduler = lr_scheduler_s2s[cur_area_idx]
optimizer = optimizer_s2s[cur_area_idx]
else:
lr_scheduler = lr_scheduler_s3s[cur_area_idx]
optimizer = optimizer_s3s[cur_area_idx]
#if not next(net.module.fpn_net.parameters()).is_cuda:
# net.module.fpn_net.cuda()
#test_data_loader = dis_test_data_loaders[data_loader_idx]
#train_data_loader = DataLoader(dis_train_sets[data_loader_idx], batch_size=48,
# shuffle=True, num_workers=2)
#test_data_loaders = DataLoader(dis_test_sets[data_loader_idx], batch_size=32,
# shuffle=False, num_workers=2)
for i, data in tqdm(enumerate(train_data_loader)):
image, face_image, gaze_field, eye_position, gt_position, gt_heatmap = \
data['image'], data['face_image'], data['gaze_field'], data['eye_position'], data['gt_position'], data['gt_heatmap']
image, face_image, gaze_field, eye_position, gt_position, gt_heatmap = \
map(lambda x: x.cuda(), [image, face_image, gaze_field, eye_position, gt_position, gt_heatmap])
# for var in [image, face_image, gaze_field, eye_position, gt_position]:
# print var.shape
optimizer.zero_grad()
direction, predict_heatmap = net(
[image, face_image, gaze_field, eye_position])
heatmap_loss, m_angle_loss = \
F_loss(direction, predict_heatmap, eye_position, gt_position, gt_heatmap)
if epoch == 0:
#if epoch < 7:
loss = m_angle_loss
elif epoch >= 7 and epoch <= 14:
loss = heatmap_loss
else:
loss = m_angle_loss + heatmap_loss
loss.backward()
optimizer.step()
# running_loss.append([heatmap_loss.data[0],
# m_angle_loss.data[0], loss.data[0]])
running_loss.append(
[heatmap_loss.item(),
m_angle_loss.item(),
loss.item()])
if i % 10 == 9:
logging.info('%s %s %s' %
(str(np.mean(running_loss, axis=0)), method,
str(lr_scheduler.get_last_lr())))
running_loss = []
lr_scheduler.step()
epoch += 1
save_path = '../model/two_stage_fpn_concat_multi_scale_' + method
if not os.path.exists(save_path):
os.makedirs(save_path)
if epoch % 5 == 0:
torch.save(net.state_dict(),
save_path + '/model_epoch{}.pkl'.format(epoch))
for i in range(16):
torch.save(net.module.fpn_nets[i].state_dict(),
save_path + '/fpn_{}.pkl'.format(i))
for data_loader_idx in range(len(dis_test_data_loaders)):
test_data_loader = dis_test_data_loaders[data_loader_idx]
if epoch > 10:
area_idx = int(data_loader_idx / area_in_network)
net.module.change_fpn(area_idx)
cur_area_idx = area_idx
test(net, test_data_loader)