| 8 | from collections import deque |
| 9 | |
| 10 | class DataLoader: |
| 11 | |
| 12 | def __init__( |
| 13 | self, |
| 14 | data_path, |
| 15 | data_file, |
| 16 | batch_size, |
| 17 | data_file_num, |
| 18 | sleep_time=1, |
| 19 | max_queue_size = 2 |
| 20 | ): |
| 21 | # load data |
| 22 | self.queue = deque() #multiprocessing.Queue(maxsize=max_queue_size) # it may change in future if we decide to split data into many small chunks instead of 4 |
| 23 | self.batch_size = batch_size |
| 24 | self.data_path = data_path |
| 25 | self.data_file = data_file |
| 26 | self.data_file_num = data_file_num |
| 27 | self.sleep_time = sleep_time |
| 28 | self.max_queue_size = max_queue_size |
| 29 | self.help_count = 0 |
| 30 | |
| 31 | def __iter__(self): |
| 32 | return self |
| 33 | |
| 34 | def data_read(self, start_id, total_thread): |
| 35 | sample_id = start_id |
| 36 | while sample_id < self.data_file_num: |
| 37 | if len(self.queue) >= self.max_queue_size: |
| 38 | time.sleep(1) |
| 39 | continue |
| 40 | processed_data_path = self.data_path + self.data_file + "_" + str(sample_id) + '_processed.npz' |
| 41 | print('Start loading processed data...' + processed_data_path) |
| 42 | st = time.time() |
| 43 | data = np.load(processed_data_path) |
| 44 | source = data['source_array'] |
| 45 | uid_array = np.array(source)[:,0] |
| 46 | item_array = np.array(source)[:,1] |
| 47 | cate_array = np.array(source)[:,2] |
| 48 | shop_array = np.array(source)[:,3] |
| 49 | node_array = np.array(source)[:,4] |
| 50 | product_array = np.array(source)[:,5] |
| 51 | brand_array = np.array(source)[:,6] |
| 52 | |
| 53 | target = data['target_array'] |
| 54 | history_item = data['history_item_array'] |
| 55 | history_cate = data['history_cate_array'] |
| 56 | history_shop = data['history_shop_array'] |
| 57 | history_node = data['history_node_array'] |
| 58 | history_product = data['history_product_array'] |
| 59 | history_brand = data['history_brand_array'] |
| 60 | |
| 61 | neg_history_item = data['neg_history_item_array'] |
| 62 | neg_history_cate = data['neg_history_cate_array'] |
| 63 | neg_history_shop = data['neg_history_shop_array'] |
| 64 | neg_history_node = data['neg_history_node_array'] |
| 65 | neg_history_product = data['neg_history_product_array'] |
| 66 | neg_history_brand = data['neg_history_brand_array'] |
| 67 | print('Finish loading processed data id '+ str(sample_id) + ',Time cost = %.4f' % (time.time()-st)) |