-
Notifications
You must be signed in to change notification settings - Fork 3
/
Copy pathutil.py
96 lines (72 loc) · 2.98 KB
/
util.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
import torch
from scipy.io import loadmat, savemat
import numpy as np
def gen_A(num_classes, t, adj_file):
result = loadmat(adj_file)
_adj = result['adj']
_nums = result['nums']
_nums = _nums[:, None]
_adj = _adj / _nums
_adj[_adj < t] = 0
_adj[_adj >= t] = 1
_adj = _adj.squeeze()
_adj = _adj * 0.5
_adj = _adj + np.identity(num_classes, np.float)
return _adj
def gen_adj(A):
D = torch.pow(A.sum(axis=1).float(), -0.5)
D = torch.diag(D)
adj = torch.matmul(torch.matmul(A, D).t(), D)
return adj
"""
#based on http://stackoverflow.com/questions/7323664/python-generator-pre-fetch
"""
import threading
import sys
if sys.version_info >= (3, 0):
import queue as Queue
else:
import Queue
class BackgroundGenerator(threading.Thread):
def __init__(self, generator, max_prefetch=1):
"""
This function transforms generator into a background-thead generator.
:param generator: generator or genexp or any
It can be used with any minibatch generator.
It is quite lightweight, but not entirely weightless.
Using global variables inside generator is not recommended (may rise GIL and zero-out the benefit of having a background thread.)
The ideal use case is when everything it requires is store inside it and everything it outputs is passed through queue.
There's no restriction on doing weird stuff, reading/writing files, retrieving URLs [or whatever] wlilst iterating.
:param max_prefetch: defines, how many iterations (at most) can background generator keep stored at any moment of time.
Whenever there's already max_prefetch batches stored in queue, the background process will halt until one of these batches is dequeued.
!Default max_prefetch=1 is okay unless you deal with some weird file IO in your generator!
Setting max_prefetch to -1 lets it store as many batches as it can, which will work slightly (if any) faster, but will require storing
all batches in memory. If you use infinite generator with max_prefetch=-1, it will exceed the RAM size unless dequeued quickly enough.
"""
threading.Thread.__init__(self)
self.queue = Queue.Queue(max_prefetch)
self.generator = generator
self.daemon = True
self.start()
def run(self):
for item in self.generator:
self.queue.put(item)
self.queue.put(None)
def next(self):
next_item = self.queue.get()
if next_item is None:
raise StopIteration
return next_item
# Python 3 compatibility
def __next__(self):
return self.next()
def __iter__(self):
return self
# decorator
class background:
def __init__(self, max_prefetch=1):
self.max_prefetch = max_prefetch
def __call__(self, gen):
def bg_generator(*args, **kwargs):
return BackgroundGenerator(gen(*args, **kwargs), max_prefetch=self.max_prefetch)
return bg_generator