-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
121 lines (67 loc) · 2.79 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
import torch
import train
import test
import numpy as np
from scipy.io import loadmat
import multi_modal_nn as mmnn
from face_landmark_dataset import FaceLandmarksDataset
import sys
import os
def extract_data(indices, suffix = '', include_pos = True):
ftrs = torch.Tensor()
ftrs = ftrs.type(torch.cuda.FloatTensor)
gz = torch.Tensor()
gz = gz.type(torch.cuda.FloatTensor)
eye_reg = torch.Tensor()
eye_reg = eye_reg.type(torch.cuda.IntTensor)
img_loc = np.asarray([])
# img_loc = img_loc.type(torch.cuda.FloatTensor)
for index in indices:
data = loadmat(str(index) + suffix + '_lmarks_location_eye.mat')
#Landmark features
ftrs_single = torch.from_numpy(data['ftrs'])
ftrs_single = ftrs_single.type(torch.cuda.FloatTensor)
ftrs = torch.cat((ftrs, ftrs_single))
#Gaze features
gz_single = torch.from_numpy(data['gz'])
gz_single = gz_single.type(torch.cuda.FloatTensor)
gz_single = torch.t(gz_single)
gz = torch.cat((gz, gz_single))
#Eye regions should be n X 4 size
eye_reg_single = torch.from_numpy(data['eye_reg'])
eye_reg_single = eye_reg_single.type(torch.cuda.IntTensor)
eye_reg = torch.cat((eye_reg, eye_reg_single))
#Get image location
img_loc_single = data['location']
img_loc = np.concatenate((img_loc, img_loc_single))
return(ftrs, gz, eye_reg, img_loc)
device = torch.device("cuda:0")
train_indices = [401,402,403,405]
# test_indices = [404, 407,410]
test_indices = [404,407,410]
ftrs = torch.Tensor()
gz = torch.Tensor()
(ftrs, gz, eye_reg, img_loc) = extract_data(train_indices)
test_ftrs = torch.Tensor().cuda()
test_gz = torch.Tensor().cuda()
(test_ftrs, test_gz, test_eye_reg, test_img_loc) = extract_data(test_indices)
#----------------------------------------------
(net,optimizer) = mmnn.get_net_instance()
face_landmarks_dataset = FaceLandmarksDataset(ftrs = ftrs, eye_regions= eye_reg.cuda(), locations=img_loc, gz = gz.cuda(), train_transforms=None, test_transforms=None)
print(torch.cuda.memory_allocated())
y = gz.cuda()
# x.requires_grad = False
y.requires_grad = False
# x = train_input
# y = train_output
id = int(sys.argv[1])
batch_size = int(sys.argv[2])
try:
# os.mkdir('log_results/%d' % (id))
os.mkdir('log_results/%d/models' % (id))
except OSError:
print ("Creation of the directory %s failed" % path)
output_file = 'log_results/%d/multimodal_log_%d_%d.txt' % (id, id, batch_size)
output_model_file = 'log_results/%d/models/multimodal_model_%d_%d' % (id, id, batch_size)
net = train.train_model(optimizer,output_model_file, net, face_landmarks_dataset , 20, 1, False, batch_size = batch_size, output_file = output_file, train_id = id)
torch.save(net.state_dict(), output_model_file)