-
Notifications
You must be signed in to change notification settings - Fork 3
/
S2_svd.py
67 lines (52 loc) · 2.14 KB
/
S2_svd.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
from libs.utils import *
class SVD(object):
def __init__(self, cfg, dict_DB):
self.cfg = cfg
self.dataloader = dict_DB['dataloader']
self.visualize = dict_DB['visualize']
self.height = cfg.height
self.width = cfg.width
self.size = np.float32([cfg.height, cfg.width])
self.center = np.array([(self.width - 1) / 2, (self.height - 1) / 2])
self.datalist = []
def update_matrix(self, data):
self.mat = torch.cat((self.mat, data.view(-1, 1)), dim=1)
def load_contour_component(self):
datalist = load_pickle(self.cfg.output_dir + 'pickle/datalist')
# sampled
for i in range(len(datalist)):
img_name = datalist[i]
data_p = load_pickle(self.cfg.output_dir + 'pickle/' + img_name)[0]
for k in range(len(data_p)):
if len(data_p[k]['r']) != 0:
self.update_matrix(torch.tensor(data_p[k]['r']).type(torch.float32).cuda())
print('%d done!' % i)
if self.cfg.save_pickle == True:
save_pickle(dir_name=self.cfg.output_dir,
file_name='matrix',
data=self.mat)
def do_SVD(self):
self.mat = load_pickle(self.cfg.output_dir + 'matrix')
n, l = self.mat.shape
idx = torch.linspace(0, l-1, 40000).type(torch.int64).cuda()
U, S, V = torch.svd(self.mat[:, idx].cpu() / (self.cfg.max_dist))
self.U = U.cuda()
self.S = S.cuda()
self.V = V.cuda()
if self.cfg.save_pickle == True:
save_pickle(dir_name=self.cfg.output_dir,
file_name='U',
data=self.U)
save_pickle(dir_name=self.cfg.output_dir,
file_name='S',
data=self.S)
def make_dict(self):
self.mat = torch.FloatTensor([]).cuda()
self.U = torch.FloatTensor([]).cuda()
self.S = torch.FloatTensor([]).cuda()
self.V = torch.FloatTensor([]).cuda()
def run(self):
print('start')
self.make_dict()
self.load_contour_component()
self.do_SVD()