E_model.cuda()
loss_f_ARE = loss_func.loss_f_ARE()
loss_f_ARE.cuda()
loss_f_E = loss_func.loss_f_E()
loss_f_E.cuda()
#img_train_list, img_test_list = dataset.load_all_h5(H5_address)
txt_list = os.listdir(H5_address)
txt_list.sort(key=lambda x: int(x[1:3]))
set_dict = dict(zip(person_set, txt_list))
train_list = []
test_list = []
for key in set_dict:
h5_list = dataset.load_h5_list(os.path.join(H5_address, set_dict[key]))
if key != person_num:
train_list.extend(h5_list)
else:
test_list.extend(h5_list)
train_Dataset = dataset.gaze_dataset(train_list)
test_Dataset = dataset.gaze_dataset(test_list)
train_loader = torch.utils.data.DataLoader(train_Dataset,
shuffle=True,
batch_size=BatchSize,
num_workers=4)
test_loader = torch.utils.data.DataLoader(test_Dataset,
shuffle=True,
batch_size=BatchSize,
num_workers=4)
#AR_down_model = nn.DataParallel(AR_down_model)
#AR_down_model.cuda()
#AR_up_model = model_ns.AR_Net_up()
#AR_up_model = nn.DataParallel(AR_down_model)
#AR_up_model.cuda()
loss_f_ARE = loss_func.loss_f_ARE()
loss_f_ARE.cuda()
loss_f_E = loss_func.loss_f_E()
loss_f_E.cuda()
#img_train_list, img_test_list = dataset.load_all_h5(H5_address)
train_list = dataset.load_all_h5(H5_train_address)
test_list = dataset.load_h5_list(H5_test_address)
train_Dataset = dataset.gaze_train_dataset(train_list)
test_Dataset = dataset.gaze_test_dataset(test_list)
train_loader = torch.utils.data.DataLoader(train_Dataset,
shuffle=True,
batch_size=BatchSize,
num_workers=4)
test_loader = torch.utils.data.DataLoader(test_Dataset,
shuffle=True,
batch_size=BatchSize,
num_workers=4)
#L1_loss = nn.SmoothL1Loss().cuda()
#l1_loss = nn.MSELoss().cuda()
#optimizer = torch.optim.Adam(gaze_model.parameters(),lr=0.01)
