test=develop, fix

Author: 123malin

models/treebased/builder/tree_index_builder.py

@@ -16,11 +16,18 @@
 import numpy as np
 import struct
 import argparse
+import os
+import time
+import collections
+import multiprocessing as mp
+
+from sklearn.cluster import KMeans
 
 
 class TreeIndexBuilder:
-    def __init__(self, branch=2):
-        self.branch = branch
+    def __init__(self):
+        self.branch = 2
+        self.timeout = 5
 
     def build_by_category(self, input_filename, output_filename):
         class Item:
@@ -64,20 +71,159 @@ def gen_code(start, end, code):
                 gen_code(start, _sub_end, self.branch * code + self.branch - i)
                 start = _sub_end
 
-            # mid = int((start + end) / 2)
-            # gen_code(mid, end, 2 * code + 1)
-            # gen_code(start, mid, 2 * code + 2)
-
         gen_code(0, len(items), 0)
         ids = np.array([item.item_id for item in items])
         codes = np.array([item.code for item in items])
-        # for i in range(len(items)):
-        #     print(ids[i], codes[i])
-        #data = np.array([[] for i in range(len(ids))])
         self.build(output_filename, ids, codes)
 
-    def tree_init_by_kmeans(self):
-        pass
+    def tree_init_by_kmeans(self, input_filename, output_filename, parall=1):
+        t1 = time.time()
+        ids = list()
+        data = list()
+        with open(input_filename) as f:
+            for line in f:
+                arr = line.split(',')
+                if not arr:
+                    break
+                ids.append(int(arr[0]))
+                vector = list()
+                for i in range(1, len(arr)):
+                    vector.append(float(arr[i]))
+                data.append(vector)
+        self.ids = np.array(ids)
+        self.data = np.array(data)
+        t2 = time.time()
+        print("Read data done, {} records read, elapsed: {}".format(
+            len(ids), t2 - t1))
+
+        queue = mp.Queue()
+        queue.put((0, np.array(range(len(self.ids)))))
+        processes = []
+        pipes = []
+        for _ in range(parall):
+            a, b = mp.Pipe()
+            p = mp.Process(target=self._train, args=(b, queue))
+            processes.append(p)
+            pipes.append(a)
+            p.start()
+
+        self.codes = np.zeros((len(self.ids), ), dtype=np.int64)
+        for pipe in pipes:
+            codes = pipe.recv()
+            for i in range(len(codes)):
+                if codes[i] > 0:
+                    self.codes[i] = codes[i]
+
+        for p in processes:
+            p.join()
+
+        assert (queue.empty())
+        self.build(output_filename, self.ids, self.codes, data=self.data)
+
+    def _train(self, pipe, queue):
+        last_size = -1
+        catch_time = 0
+        processed = False
+        code = np.zeros((len(self.ids), ), dtype=np.int64)
+        while True:
+            for _ in range(5):
+                try:
+                    pcode, index = queue.get(timeout=self.timeout)
+                except:
+                    index = None
+                if index is not None:
+                    break
+
+            if index is None:
+                if processed and (last_size <= 1024 or catch_time >= 3):
+                    print("Process {} exits".format(os.getpid()))
+                    break
+                else:
+                    print("Got empty job, pid: {}, time: {}".format(os.getpid(
+                    ), catch_time))
+                    catch_time += 1
+                    continue
+
+            processed = True
+            catch_time = 0
+            last_size = len(index)
+            if last_size <= 1024:
+                self._minbatch(pcode, index, code)
+            else:
+                tstart = time.time()
+                left_index, right_index = self._cluster(index)
+                if last_size > 1024:
+                    print("Train iteration done, pcode:{}, "
+                          "data size: {}, elapsed time: {}"
+                          .format(pcode, len(index), time.time() - tstart))
+                self.timeout = int(0.4 * self.timeout + 0.6 * (time.time() -
+                                                               tstart))
+                if self.timeout < 5:
+                    self.timeout = 5
+
+                if len(left_index) > 1:
+                    queue.put((2 * pcode + 1, left_index))
+
+                if len(right_index) > 1:
+                    queue.put((2 * pcode + 2, right_index))
+        process_count = 0
+        for c in code:
+            if c > 0:
+                process_count += 1
+        print("Process {} process {} items".format(os.getpid(), process_count))
+        pipe.send(code)
+
+    def _minbatch(self, pcode, index, code):
+        dq = collections.deque()
+        dq.append((pcode, index))
+        batch_size = len(index)
+        tstart = time.time()
+        while dq:
+            pcode, index = dq.popleft()
+
+            if len(index) == 2:
+                code[index[0]] = 2 * pcode + 1
+                code[index[1]] = 2 * pcode + 2
+                continue
+
+            left_index, right_index = self._cluster(index)
+            if len(left_index) > 1:
+                dq.append((2 * pcode + 1, left_index))
+            elif len(left_index) == 1:
+                code[left_index] = 2 * pcode + 1
+
+            if len(right_index) > 1:
+                dq.append((2 * pcode + 2, right_index))
+            elif len(right_index) == 1:
+                code[right_index] = 2 * pcode + 2
+
+        print("Minbatch, batch size: {}, elapsed: {}".format(
+            batch_size, time.time() - tstart))
+
+    def _cluster(self, index):
+        data = self.data[index]
+        kmeans = KMeans(n_clusters=2, random_state=0).fit(data)
+        labels = kmeans.labels_
+        l_i = np.where(labels == 0)[0]
+        r_i = np.where(labels == 1)[0]
+        left_index = index[l_i]
+        right_index = index[r_i]
+        if len(right_index) - len(left_index) > 1:
+            distances = kmeans.transform(data[r_i])
+            left_index, right_index = self._rebalance(left_index, right_index,
+                                                      distances[:, 1])
+        elif len(left_index) - len(right_index) > 1:
+            distances = kmeans.transform(data[l_i])
+            left_index, right_index = self._rebalance(right_index, left_index,
+                                                      distances[:, 0])
+
+        return left_index, right_index
+
+    def _rebalance(self, lindex, rindex, distances):
+        sorted_index = rindex[np.argsort(distances)]
+        idx = np.concatenate((lindex, sorted_index))
+        mid = int(len(idx) / 2)
+        return idx[mid:], idx[:mid]
 
     def build(self, output_filename, ids, codes, data=None, id_offset=None):
         # process id offset
@@ -161,7 +307,11 @@ def _write_kv(self, fwr, message):
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description="TreeIndexBuiler")
     parser.add_argument(
-        "--branch", required=False, type=int, default=2, help="tree branch.")
+        "--parallel",
+        required=False,
+        type=int,
+        default=12,
+        help="parallel nums.")
     parser.add_argument(
         "--mode",
         required=True,
@@ -172,8 +322,8 @@ def _write_kv(self, fwr, message):
 
     args = parser.parse_args()
     if args.mode == "by_category":
-        builder = TreeIndexBuilder(args.branch)
+        builder = TreeIndexBuilder()
         builder.build_by_category(args.input, args.output)
     elif args.mode == "by_kmeans":
-        builder = TreeIndexBuilder(args.branch)
-        builder.tree_init_by_category(args.input, args.output)
+        builder = TreeIndexBuilder()
+        builder.tree_init_by_kmeans(args.input, args.output, args.parallel)

models/treebased/tdm/config.yaml

@@ -31,7 +31,7 @@ runner:
   train_batch_size: 100  # 30000
   epochs: 5
   print_interval: 10   # 1000
-  model_save_path: "tdm_demo_output"
+  model_save_path: "output_model_tdm_demo"
 
 # hyper parameters of user-defined network
 hyper_parameters:

models/treebased/tdm/config_ub.yaml

@@ -20,17 +20,17 @@ runner:
 
   model_path: "static_model.py"
   reader_type: "QueueDataset"  # DataLoader / QueueDataset / RecDataset
-  pipe_command: "python ub_reader.py"
+  pipe_command: "python reader.py"
   dataset_debug: False
   split_file_list: False
 
-  train_data_dir: "../ub_data/debug_data"
-  train_reader_path: "ub_reader"
+  train_data_dir: "../ub_data/train_data"
+  train_reader_path: "reader"
 
-  train_batch_size: 300
+  train_batch_size: 30000
   epochs: 5
-  print_interval: 10   # 1000
-  model_save_path: "tdm_demo_ub"
+  print_interval: 1000
+  model_save_path: "output_model_tdm_ub"
 
 # hyper parameters of user-defined network
 hyper_parameters: