forked from uoguelph-mlrg/theano_alexnet
-
Notifications
You must be signed in to change notification settings - Fork 0
/
proc_load.py
143 lines (106 loc) · 4.27 KB
/
proc_load.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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
'''
Load data in parallel with train.py
'''
import time
import math
import numpy as np
import zmq
import pycuda.driver as drv
import pycuda.gpuarray as gpuarray
import hickle as hkl
import pickle
def get_params_crop_and_mirror(param_rand, data_shape, cropsize):
center_margin = (data_shape[2] - cropsize) / 2
crop_xs = round(param_rand[0] * center_margin * 2)
crop_ys = round(param_rand[1] * center_margin * 2)
if False:
# this is true then exactly replicate Ryan's code, in the batch case
crop_xs = math.floor(param_rand[0] * center_margin * 2)
crop_ys = math.floor(param_rand[1] * center_margin * 2)
flag_mirror = bool(round(param_rand[2]))
return crop_xs, crop_ys, flag_mirror
def crop_and_mirror(data, param_rand, flag_batch=True, cropsize=227):
'''
when param_rand == (0.5, 0.5, 0), it means no randomness
'''
# print param_rand
# if param_rand == (0.5, 0.5, 0), means no randomness and do validation
if param_rand[0] == 0.5 and param_rand[1] == 0.5 and param_rand[2] == 0:
flag_batch = True
if flag_batch:
# mirror and crop the whole batch
crop_xs, crop_ys, flag_mirror = \
get_params_crop_and_mirror(param_rand, data.shape, cropsize)
# random mirror
if flag_mirror:
data = data[:, :, ::-1, :]
# random crop
data = data[:, crop_xs:crop_xs + cropsize,
crop_ys:crop_ys + cropsize, :]
else:
# mirror and crop each batch individually
# to ensure consistency, use the param_rand[1] as seed
np.random.seed(int(10000 * param_rand[1]))
data_out = np.zeros((data.shape[0], cropsize, cropsize,
data.shape[3])).astype('float32')
for ind in range(data.shape[3]):
# generate random numbers
tmp_rand = np.float32(np.random.rand(3))
tmp_rand[2] = round(tmp_rand[2])
# get mirror/crop parameters
crop_xs, crop_ys, flag_mirror = \
get_params_crop_and_mirror(tmp_rand, data.shape, cropsize)
# do image crop/mirror
img = data[:, :, :, ind]
if flag_mirror:
img = img[:, :, ::-1]
img = img[:, crop_xs:crop_xs + cropsize,
crop_ys:crop_ys + cropsize]
data_out[:, :, :, ind] = img
data = data_out
return np.ascontiguousarray(data, dtype='float32')
def fun_load(config, sock_data=5000):
send_queue = config['queue_l2t']
recv_queue = config['queue_t2l']
# recv_queue and send_queue are multiprocessing.Queue
# recv_queue is only for receiving
# send_queue is only for sending
# if need to do random crop and mirror
flag_randproc = not config['use_data_layer']
flag_batch = config['batch_crop_mirror']
drv.init()
dev = drv.Device(int(config['gpu'][-1]))
ctx = dev.make_context()
sock = zmq.Context().socket(zmq.PAIR)
sock.bind('tcp://*:{0}'.format(sock_data))
shape, dtype, h = sock.recv_pyobj()
print 'shared_x information received', shape, dtype
shape = (3, 255, 255, 256) # TODO remove fix
gpu_data_remote = gpuarray.GPUArray(shape, dtype,
gpudata=drv.IPCMemoryHandle(h))
gpu_data = gpuarray.GPUArray(shape, dtype)
img_mean = recv_queue.get()
print 'img_mean received'
# The first time, do the set ups and other stuff
# receive information for loading
while True:
# getting the hkl file name to load
hkl_name = recv_queue.get()
# print hkl_name
#data = pickle.load(open(hkl_name)) - img_mean
data = hkl.load(hkl_name) - img_mean
# print 'load ', time.time() - bgn_time
if flag_randproc:
param_rand = recv_queue.get()
data = crop_and_mirror(data, param_rand, flag_batch=flag_batch)
gpu_data.set(data)
# wait for computation on last minibatch to finish
msg = recv_queue.get()
assert msg == 'calc_finished'
drv.memcpy_peer(gpu_data_remote.ptr,
gpu_data.ptr,
gpu_data.dtype.itemsize *
gpu_data.size,
ctx, ctx)
ctx.synchronize()
send_queue.put('copy_finished')