easy-cs-rec-custommodel 0.8.6__py2.py3-none-any.whl

easy_rec/python/inference/processor/test.py

@@ -0,0 +1,170 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import argparse
+import ctypes
+import glob
+import json
+import logging
+import os
+import subprocess
+import time
+
+import numpy as np
+from google.protobuf import text_format
+
+from easy_rec.python.protos import dataset_pb2
+from easy_rec.python.protos import pipeline_pb2
+from easy_rec.python.protos import tf_predict_pb2
+
+logging.basicConfig(
+    level=logging.INFO, format='[%(asctime)s][%(levelname)s] %(message)s')
+
+PROCESSOR_VERSION = 'LaRec-0.9.5d-b1b1604-TF-2.5.0-Linux'
+PROCESSOR_FILE = PROCESSOR_VERSION + '.tar.gz'
+PROCESSOR_URL = 'http://easyrec.oss-cn-beijing.aliyuncs.com/processor/' + PROCESSOR_FILE
+PROCESSOR_ENTRY_LIB = 'processor/' + PROCESSOR_VERSION + '/larec/libtf_predictor.so'
+
+
+def build_array_proto(array_proto, data, dtype):
+  array_proto.array_shape.dim.append(len(data))
+
+  if dtype == dataset_pb2.DatasetConfig.STRING:
+    array_proto.string_val.extend([x.encode('utf-8') for x in data])
+    array_proto.dtype = tf_predict_pb2.DT_STRING
+  elif dtype == dataset_pb2.DatasetConfig.FLOAT:
+    array_proto.float_val.extend([float(x) for x in data])
+    array_proto.dtype = tf_predict_pb2.DT_FLOAT
+  elif dtype == dataset_pb2.DatasetConfig.DOUBLE:
+    array_proto.double_val.extend([float(x) for x in data])
+    array_proto.dtype = tf_predict_pb2.DT_DOUBLE
+  elif dtype == dataset_pb2.DatasetConfig.INT32:
+    array_proto.int_val.extend([int(x) for x in data])
+    array_proto.dtype = tf_predict_pb2.DT_INT32
+  elif dtype == dataset_pb2.DatasetConfig.INT64:
+    array_proto.int64_val.extend([np.int64(x) for x in data])
+    array_proto.dtype = tf_predict_pb2.DT_INT64
+  else:
+    assert False, 'invalid datatype[%s]' % str(dtype)
+  return array_proto
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+      '--input_path', type=str, default=None, help='input data path')
+  parser.add_argument(
+      '--output_path', type=str, default=None, help='output data path')
+  parser.add_argument(
+      '--libc_path',
+      type=str,
+      default='/lib64/libc.so.6',
+      help='libc.so.6 path')
+  parser.add_argument(
+      '--saved_model_dir', type=str, default=None, help='saved model directory')
+  parser.add_argument(
+      '--test_dir', type=str, default=None, help='test directory')
+  args = parser.parse_args()
+
+  if not os.path.exists('processor'):
+    os.mkdir('processor')
+  if not os.path.exists(PROCESSOR_ENTRY_LIB):
+    if not os.path.exists('processor/' + PROCESSOR_FILE):
+      subprocess.check_output(
+          'wget %s -O processor/%s' % (PROCESSOR_URL, PROCESSOR_FILE),
+          shell=True)
+    subprocess.check_output(
+        'cd processor && tar -zvxf %s' % PROCESSOR_FILE, shell=True)
+    assert os.path.exists(
+        PROCESSOR_ENTRY_LIB), 'invalid processor path: %s' % PROCESSOR_ENTRY_LIB
+
+  assert os.path.exists(args.libc_path), '%s does not exist' % args.libc_path
+  assert args.saved_model_dir is not None and os.path.isdir(
+      args.saved_model_dir
+  ), '%s is not a valid directory' % args.saved_model_dir
+  assert args.input_path is not None and os.path.exists(
+      args.input_path), '%s does not exist' % args.input_path
+  assert args.output_path is not None, 'output_path is not set'
+
+  pipeline_config = pipeline_pb2.EasyRecConfig()
+  pipeline_config_path = os.path.join(args.saved_model_dir,
+                                      'assets/pipeline.config')
+  with open(pipeline_config_path) as fin:
+    config_str = fin.read()
+  text_format.Merge(config_str, pipeline_config)
+
+  data_config = pipeline_config.data_config
+
+  input_fields = [[]
+                  for x in data_config.input_fields
+                  if x.input_name not in data_config.label_fields]
+
+  with open(args.input_path, 'r') as fin:
+    for line_str in fin:
+      line_str = line_str.strip()
+      line_toks = line_str.split(data_config.rtp_separator)[-1].split(chr(2))
+      for i, tok in enumerate(line_toks):
+        input_fields[i].append(tok)
+
+  req = tf_predict_pb2.PredictRequest()
+  req.signature_name = 'serving_default'
+  for i in range(len(input_fields)):
+    build_array_proto(req.inputs[data_config.input_fields[i + 1].input_name],
+                      input_fields[i],
+                      data_config.input_fields[i + 1].input_type)
+
+  tf_predictor = ctypes.cdll.LoadLibrary(PROCESSOR_ENTRY_LIB)
+  tf_predictor.saved_model_init.restype = ctypes.c_void_p
+  handle = tf_predictor.saved_model_init(args.saved_model_dir.encode('utf-8'))
+  logging.info('saved_model handle=%d' % handle)
+
+  num_steps = pipeline_config.train_config.num_steps
+  logging.info('num_steps=%d' % num_steps)
+
+  # last_step could be greater than num_steps for sync_replicas: false
+  train_dir = os.path.dirname(args.saved_model_dir.strip('/'))
+  all_models = glob.glob(
+      os.path.join(args.test_dir, 'train/model.ckpt-*.index'))
+  iters = [int(x.split('-')[-1].replace('.index', '')) for x in all_models]
+  iters.sort()
+  last_step = iters[-1]
+  logging.info('last_step=%d' % last_step)
+
+  sparse_step = ctypes.c_int(0)
+  dense_step = ctypes.c_int(0)
+  start_ts = time.time()
+  while sparse_step.value < last_step or dense_step.value < last_step:
+    tf_predictor.saved_model_step(
+        ctypes.c_void_p(handle), ctypes.byref(sparse_step),
+        ctypes.byref(dense_step))
+    time.sleep(1)
+    if time.time() - start_ts > 300:
+      logging.warning(
+          'could not reach last_step, sparse_step=%d dense_step=%d' %
+          (sparse_step.value, dense_step.value))
+      break
+
+  data_bin = req.SerializeToString()
+  save_path = os.path.join(args.saved_model_dir, 'req.pb')
+  with open(save_path, 'wb') as fout:
+    fout.write(data_bin)
+  logging.info('save request to %s' % save_path)
+
+  tf_predictor.saved_model_predict.restype = ctypes.c_void_p
+  out_len = ctypes.c_int(0)
+  res_p = tf_predictor.saved_model_predict(
+      ctypes.c_void_p(handle), data_bin, ctypes.c_int32(len(data_bin)),
+      ctypes.byref(out_len))
+  res_bytes = bytearray(ctypes.string_at(res_p, out_len))
+  res = tf_predict_pb2.PredictResponse()
+  res.ParseFromString(res_bytes)
+
+  with open(args.output_path, 'w') as fout:
+    logits = res.outputs['logits'].float_val
+    probs = res.outputs['probs'].float_val
+    for logit, prob in zip(logits, probs):
+      fout.write(json.dumps({'logits': logit, 'probs': prob}) + '\n')
+
+  # free memory
+  tf_predictor.saved_model_release(ctypes.c_void_p(handle))
+  libc = ctypes.cdll.LoadLibrary(args.libc_path)
+  libc.free(ctypes.c_void_p(res_p))

easy_rec/python/inference/vector_retrieve.py

@@ -0,0 +1,124 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import logging
+from datetime import datetime
+
+import common_io
+import numpy as np
+import tensorflow as tf
+
+try:
+  import graphlearn as gl
+except:  # noqa: E722
+  logging.warning(
+      'GraphLearn is not installed. You can install it by "pip install https://easyrec.oss-cn-beijing.aliyuncs.com/3rdparty/graphlearn-0.7-cp27-cp27mu-linux_x86_64.whl.'  # noqa: E501
+  )
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+class VectorRetrieve(object):
+
+  def __init__(self,
+               query_table,
+               doc_table,
+               out_table,
+               ndim,
+               delimiter=',',
+               batch_size=4,
+               index_type='ivfflat',
+               nlist=10,
+               nprobe=2,
+               distance=1,
+               m=8):
+    """Retrieve top n neighbours by query vector.
+
+    Args:
+      query_table: query vector table
+      doc_table: document vector table
+      out_table: output table
+      ndim: int, number of feature dimensions
+      delimiter: delimiter for feature vectors
+      batch_size: query batch size
+      index_type: search model `flat`, `ivfflat`, `ivfpq`, `gpu_ivfflat`
+      nlist: number of split part on each worker
+      nprobe: probe part on each worker
+      distance: type of distance,  0 is l2 distance(default), 1 is inner product.
+      m: number of dimensions for each node after compress
+    """
+    self.query_table = query_table
+    self.doc_table = doc_table
+    self.out_table = out_table
+    self.ndim = ndim
+    self.delimiter = delimiter
+    self.batch_size = batch_size
+
+    gl.set_inter_threadnum(8)
+    gl.set_knn_metric(distance)
+    knn_option = gl.IndexOption()
+    knn_option.name = 'knn'
+    knn_option.index_type = index_type
+    knn_option.nlist = nlist
+    knn_option.nprobe = nprobe
+    knn_option.m = m
+    self.knn_option = knn_option
+
+  def __call__(self, top_n, task_index, task_count, *args, **kwargs):
+    g = gl.Graph()
+    g.node(
+        self.doc_table,
+        'doc',
+        decoder=gl.Decoder(
+            attr_types=['float'] * self.ndim, attr_delimiter=self.delimiter),
+        option=self.knn_option)
+    g.init(task_index=task_index, task_count=task_count)
+
+    query_reader = common_io.table.TableReader(
+        self.query_table, slice_id=task_index, slice_count=task_count)
+    num_records = query_reader.get_row_count()
+    total_batch_num = num_records // self.batch_size + 1.0
+    batch_num = 0
+    print('total input records: {}'.format(query_reader.get_row_count()))
+    print('total_batch_num: {}'.format(total_batch_num))
+    print('output_table: {}'.format(self.out_table))
+
+    output_table_writer = common_io.table.TableWriter(self.out_table,
+                                                      task_index)
+    count = 0
+    while True:
+      try:
+        batch_query_nodes, batch_query_feats = zip(
+            *query_reader.read(self.batch_size, allow_smaller_final_batch=True))
+        batch_num += 1.0
+        print('{} process: {:.2f}'.format(datetime.now().time(),
+                                          batch_num / total_batch_num))
+        feats = to_np_array(batch_query_feats, self.delimiter)
+        rt_ids, rt_dists = g.search('doc', feats, gl.KnnOption(k=top_n))
+
+        for query_node, nodes, dists in zip(batch_query_nodes, rt_ids,
+                                            rt_dists):
+          query = np.array([query_node] * len(nodes), dtype='int64')
+          output_table_writer.write(
+              zip(query, nodes, dists), (0, 1, 2), allow_type_cast=False)
+          count += 1
+          if np.mod(count, 100) == 0:
+            print('write ', count, ' query nodes totally')
+      except Exception as e:
+        print(e)
+        break
+
+    print('==finished==')
+    output_table_writer.close()
+    query_reader.close()
+    g.close()
+
+
+def to_np_array(batch_query_feats, attr_delimiter):
+  return np.array(
+      [map(float, feat.split(attr_delimiter)) for feat in batch_query_feats],
+      dtype='float32')

easy_rec/python/input/batch_tfrecord_input.py

@@ -0,0 +1,117 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import logging
+
+import tensorflow as tf
+
+from easy_rec.python.input.input import Input
+from easy_rec.python.utils.tf_utils import get_tf_type
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+class BatchTFRecordInput(Input):
+  """BatchTFRecordInput use for batch read from tfrecord.
+
+  For example, there is a tfrecord which one feature(key)
+  correspond to n data(value).
+  batch_size needs to be a multiple of n.
+  """
+
+  def __init__(self,
+               data_config,
+               feature_config,
+               input_path,
+               task_index=0,
+               task_num=1,
+               check_mode=False,
+               pipeline_config=None):
+    super(BatchTFRecordInput,
+          self).__init__(data_config, feature_config, input_path, task_index,
+                         task_num, check_mode, pipeline_config)
+    assert data_config.HasField(
+        'n_data_batch_tfrecord'), 'Need to set n_data_batch_tfrecord in config.'
+    self._input_shapes = [x.input_shape for x in data_config.input_fields]
+    self.feature_desc = {}
+    for x, t, d, s in zip(self._input_fields, self._input_field_types,
+                          self._input_field_defaults, self._input_shapes):
+      d = self.get_type_defaults(t, d)
+      t = get_tf_type(t)
+      self.feature_desc[x] = tf.io.FixedLenSequenceFeature(
+          dtype=t, shape=s, allow_missing=True)
+
+  def _parse_tfrecord(self, example):
+    try:
+      _, features, _ = tf.parse_sequence_example(
+          example, sequence_features=self.feature_desc)
+    except AttributeError:
+      _, features, _ = tf.io.parse_sequence_example(
+          example, sequence_features=self.feature_desc)
+    # Below code will reduce one dimension when the data dimension > 2.
+    features = dict(
+        (key,
+         tf.reshape(value, [
+             -1,
+         ] + [x for i, x in enumerate(value.shape) if i not in (0, 1)])) for (
+             key, value) in features.items())
+    return features
+
+  def _build(self, mode, params):
+    if type(self._input_path) != list:
+      self._input_path = self._input_path.split(',')
+    file_paths = []
+    for x in self._input_path:
+      file_paths.extend(tf.gfile.Glob(x))
+    assert len(file_paths) > 0, 'match no files with %s' % self._input_path
+
+    num_parallel_calls = self._data_config.num_parallel_calls
+    data_compression_type = self._data_config.data_compression_type
+    if mode == tf.estimator.ModeKeys.TRAIN:
+      logging.info('train files[%d]: %s' %
+                   (len(file_paths), ','.join(file_paths)))
+      dataset = tf.data.Dataset.from_tensor_slices(file_paths)
+      if self._data_config.shuffle:
+        # shuffle input files
+        dataset = dataset.shuffle(len(file_paths))
+      # too many readers read the same file will cause performance issues
+      # as the same data will be read multiple times
+      parallel_num = min(num_parallel_calls, len(file_paths))
+      dataset = dataset.interleave(
+          lambda x: tf.data.TFRecordDataset(
+              x, compression_type=data_compression_type),
+          cycle_length=parallel_num,
+          num_parallel_calls=parallel_num)
+      dataset = dataset.shard(self._task_num, self._task_index)
+      if self._data_config.shuffle:
+        dataset = dataset.shuffle(
+            self._data_config.shuffle_buffer_size,
+            seed=2020,
+            reshuffle_each_iteration=True)
+      dataset = dataset.repeat(self.num_epochs)
+    else:
+      logging.info('eval files[%d]: %s' %
+                   (len(file_paths), ','.join(file_paths)))
+      dataset = tf.data.TFRecordDataset(
+          file_paths, compression_type=data_compression_type)
+      dataset = dataset.repeat(1)
+
+    # We read n data from tfrecord one time.
+    cur_batch = self._data_config.batch_size // self._data_config.n_data_batch_tfrecord
+    cur_batch = max(1, cur_batch)
+    dataset = dataset.batch(cur_batch)
+    dataset = dataset.map(
+        self._parse_tfrecord, num_parallel_calls=num_parallel_calls)
+
+    dataset = dataset.prefetch(buffer_size=self._prefetch_size)
+    dataset = dataset.map(
+        map_func=self._preprocess, num_parallel_calls=num_parallel_calls)
+
+    dataset = dataset.prefetch(buffer_size=self._prefetch_size)
+
+    if mode != tf.estimator.ModeKeys.PREDICT:
+      dataset = dataset.map(lambda x:
+                            (self._get_features(x), self._get_labels(x)))
+    else:
+      dataset = dataset.map(lambda x: (self._get_features(x)))
+    return dataset

easy_rec/python/input/criteo_binary_reader.py

@@ -0,0 +1,259 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import argparse
+import concurrent
+import concurrent.futures
+import glob
+import logging
+import os
+import queue
+import time
+
+import numpy as np
+
+
+class BinaryDataset:
+
+  def __init__(
+      self,
+      label_bins,
+      dense_bins,
+      category_bins,
+      batch_size=1,
+      drop_last=False,
+      prefetch=1,
+      global_rank=0,
+      global_size=1,
+  ):
+    total_sample_num = 0
+    self._sample_num_arr = []
+    for label_bin in label_bins:
+      sample_num = os.path.getsize(label_bin) // 4
+      total_sample_num += sample_num
+      self._sample_num_arr.append(sample_num)
+    logging.info('total number samples = %d' % total_sample_num)
+    self._total_sample_num = total_sample_num
+
+    self._batch_size = batch_size
+
+    self._compute_global_start_pos(total_sample_num, batch_size, global_rank,
+                                   global_size, drop_last)
+
+    self._label_file_arr = [None for _ in self._sample_num_arr]
+    self._dense_file_arr = [None for _ in self._sample_num_arr]
+    self._category_file_arr = [None for _ in self._sample_num_arr]
+
+    for tmp_file_id in range(self._start_file_id, self._end_file_id + 1):
+      self._label_file_arr[tmp_file_id] = os.open(label_bins[tmp_file_id],
+                                                  os.O_RDONLY)
+      self._dense_file_arr[tmp_file_id] = os.open(dense_bins[tmp_file_id],
+                                                  os.O_RDONLY)
+      self._category_file_arr[tmp_file_id] = os.open(category_bins[tmp_file_id],
+                                                     os.O_RDONLY)
+
+    self._prefetch = min(prefetch, self._num_entries)
+    self._prefetch_queue = queue.Queue()
+    self._executor = concurrent.futures.ThreadPoolExecutor(
+        max_workers=self._prefetch)
+
+    self._os_close_func = os.close
+
+  def _compute_global_start_pos(self, total_sample_num, batch_size, global_rank,
+                                global_size, drop_last):
+    # ensure all workers have the same number of samples
+    avg_sample_num = (total_sample_num // global_size)
+    res_num = (total_sample_num % global_size)
+    self._num_samples = avg_sample_num
+    if res_num > 0:
+      self._num_samples += 1
+      if global_rank < res_num:
+        global_start_pos = (avg_sample_num + 1) * global_rank
+      else:
+        global_start_pos = avg_sample_num * global_rank + res_num - 1
+    else:
+      global_start_pos = avg_sample_num * global_rank
+    # global_end_pos = global_start_pos + self._num_samples
+
+    self._num_entries = self._num_samples // batch_size
+    self._last_batch_size = batch_size
+    if not drop_last and (self._num_samples % batch_size != 0):
+      self._num_entries += 1
+      self._last_batch_size = self._num_samples % batch_size
+    logging.info('num_batches = %d num_samples = %d' %
+                 (self._num_entries, self._num_samples))
+
+    start_file_id = 0
+    curr_pos = 0
+    while curr_pos + self._sample_num_arr[start_file_id] <= global_start_pos:
+      start_file_id += 1
+      curr_pos += self._sample_num_arr[start_file_id]
+    self._start_file_id = start_file_id
+    self._start_file_pos = global_start_pos - curr_pos
+
+    logging.info('start_file_id = %d start_file_pos = %d' %
+                 (start_file_id, self._start_file_pos))
+
+    # find the start of each batch
+    self._start_pos_arr = np.zeros([self._num_entries, 2], dtype=np.uint32)
+    batch_id = 0
+    tmp_start_pos = self._start_file_pos
+    while batch_id < self._num_entries:
+      self._start_pos_arr[batch_id] = (start_file_id, tmp_start_pos)
+      batch_id += 1
+      # the last batch
+      if batch_id == self._num_entries:
+        tmp_start_pos += self._last_batch_size
+        while start_file_id < len(
+            self._sample_num_arr
+        ) and tmp_start_pos > self._sample_num_arr[start_file_id]:
+          tmp_start_pos -= self._sample_num_arr[start_file_id]
+          start_file_id += 1
+      else:
+        tmp_start_pos += batch_size
+        while start_file_id < len(
+            self._sample_num_arr
+        ) and tmp_start_pos >= self._sample_num_arr[start_file_id]:
+          tmp_start_pos -= self._sample_num_arr[start_file_id]
+          start_file_id += 1
+
+    self._end_file_id = start_file_id
+    self._end_file_pos = tmp_start_pos
+
+    logging.info('end_file_id = %d end_file_pos = %d' %
+                 (self._end_file_id, self._end_file_pos))
+
+  def __del__(self):
+    for f in self._label_file_arr:
+      if f is not None:
+        self._os_close_func(f)
+    for f in self._dense_file_arr:
+      if f is not None:
+        self._os_close_func(f)
+    for f in self._category_file_arr:
+      if f is not None:
+        self._os_close_func(f)
+
+  def __len__(self):
+    return self._num_entries
+
+  def __getitem__(self, idx):
+    if idx >= self._num_entries:
+      raise IndexError()
+
+    if self._prefetch <= 1:
+      return self._get(idx)
+
+    if idx == 0:
+      for i in range(self._prefetch):
+        self._prefetch_queue.put(self._executor.submit(self._get, (i)))
+
+    if idx < (self._num_entries - self._prefetch):
+      self._prefetch_queue.put(
+          self._executor.submit(self._get, (idx + self._prefetch)))
+
+    return self._prefetch_queue.get().result()
+
+  def _get(self, idx):
+    curr_file_id = self._start_pos_arr[idx][0]
+    start_read_pos = self._start_pos_arr[idx][1]
+
+    end_read_pos = start_read_pos + self._batch_size
+    total_read_num = 0
+
+    label_read_arr = []
+    dense_read_arr = []
+    cate_read_arr = []
+    while total_read_num < self._batch_size and curr_file_id < len(
+        self._sample_num_arr):
+      tmp_read_num = min(end_read_pos,
+                         self._sample_num_arr[curr_file_id]) - start_read_pos
+
+      label_raw_data = os.pread(self._label_file_arr[curr_file_id],
+                                4 * tmp_read_num, start_read_pos * 4)
+      tmp_lbl_np = np.frombuffer(
+          label_raw_data, dtype=np.int32).reshape([tmp_read_num, 1])
+      label_read_arr.append(tmp_lbl_np)
+
+      dense_raw_data = os.pread(self._dense_file_arr[curr_file_id],
+                                52 * tmp_read_num, start_read_pos * 52)
+      part_dense_np = np.frombuffer(
+          dense_raw_data, dtype=np.float32).reshape([tmp_read_num, 13])
+      # part_dense_np = np.log(part_dense_np + 3, dtype=np.float32)
+      dense_read_arr.append(part_dense_np)
+
+      category_raw_data = os.pread(self._category_file_arr[curr_file_id],
+                                   104 * tmp_read_num, start_read_pos * 104)
+      part_cate_np = np.frombuffer(
+          category_raw_data, dtype=np.uint32).reshape([tmp_read_num, 26])
+      cate_read_arr.append(part_cate_np)
+
+      curr_file_id += 1
+      start_read_pos = 0
+      total_read_num += tmp_read_num
+
+    if len(label_read_arr) == 1:
+      label = label_read_arr[0]
+    else:
+      label = np.concatenate(label_read_arr, axis=0)
+
+    if len(cate_read_arr) == 1:
+      category = cate_read_arr[0]
+    else:
+      category = np.concatenate(cate_read_arr, axis=0)
+
+    if len(dense_read_arr) == 1:
+      dense = dense_read_arr[0]
+    else:
+      dense = np.concatenate(dense_read_arr, axis=0)
+
+    return dense, category, label
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--batch_size', type=int, default=1024, help='batch_size')
+  parser.add_argument(
+      '--dataset_dir', type=str, default='./', help='dataset_dir')
+  parser.add_argument('--task_num', type=int, default=1, help='task number')
+  parser.add_argument('--task_index', type=int, default=0, help='task index')
+  parser.add_argument(
+      '--prefetch_size', type=int, default=10, help='prefetch size')
+  args = parser.parse_args()
+
+  batch_size = args.batch_size
+  dataset_dir = args.dataset_dir
+  logging.info('batch_size = %d' % batch_size)
+  logging.info('dataset_dir = %s' % dataset_dir)
+
+  label_files = glob.glob(os.path.join(dataset_dir, '*_label.bin'))
+  dense_files = glob.glob(os.path.join(dataset_dir, '*_dense.bin'))
+  category_files = glob.glob(os.path.join(dataset_dir, '*_category.bin'))
+
+  label_files.sort()
+  dense_files.sort()
+  category_files.sort()
+
+  test_dataset = BinaryDataset(
+      label_files,
+      dense_files,
+      category_files,
+      batch_size=batch_size,
+      drop_last=False,
+      prefetch=args.prefetch_size,
+      global_rank=args.task_index,
+      global_size=args.task_num,
+  )
+
+  for step, (dense, category, labels) in enumerate(test_dataset):
+    # if (step % 100 == 0):
+    #   print(step, dense.shape, category.shape, labels.shape)
+    if step == 0:
+      logging.info('warmup over!')
+      start_time = time.time()
+    if step == 1000:
+      logging.info('1000 steps time = %.3f' % (time.time() - start_time))
+  logging.info('total_steps = %d total_time = %.3f' %
+               (step + 1, time.time() - start_time))
+  logging.info(
+      'final step[%d] dense_shape=%s category_shape=%s labels_shape=%s' %
+      (step, dense.shape, category.shape, labels.shape))