[feature]: incremental saving checkpoint

MANIFEST.in

@@ -0,0 +1,2 @@
+include easy_rec/python/ops/1.12/*.so*
+include easy_rec/python/ops/1.15/*.so*

easy_rec/__init__.py

@@ -15,6 +15,16 @@
 logging.basicConfig(
  level=logging.INFO, format='[%(asctime)s][%(levelname)s] %(message)s')
 
+ops_dir = os.path.join(curr_dir, 'python/ops')
+if 'PAI' in tf.__version__:
+ ops_dir = os.path.join(ops_dir, '1.12_pai')
+elif tf.__version__.startswith('1.12'):
+ ops_dir = os.path.join(ops_dir, '1.12')
+elif tf.__version__.startswith('1.15'):
+ ops_dir = os.path.join(ops_dir, '1.15')
+else:
+ ops_dir = None
+
 from easy_rec.python.inference.predictor import Predictor # isort:skip # noqa: E402
 from easy_rec.python.main import evaluate # isort:skip # noqa: E402
 from easy_rec.python.main import distribute_evaluate # isort:skip # noqa: E402
@@ -32,12 +42,6 @@
 
 _global_config = {}
 
-ops_dir = os.path.join(curr_dir, 'python/ops')
-if tf.__version__.startswith('1.12'):
- ops_dir = os.path.join(ops_dir, '1.12')
-elif tf.__version__.startswith('1.15'):
- ops_dir = os.path.join(ops_dir, '1.15')
-
 
 def help():
  print("""

easy_rec/python/compat/optimizers.py

@@ -35,6 +35,9 @@
 from tensorflow.python.training import moving_averages
 from tensorflow.python.training import optimizer as optimizer_
 from tensorflow.python.training import training as train
+from easy_rec.python.ops.incr_record import set_sparse_indices
+import tensorflow as tf
+import logging
 
 OPTIMIZER_CLS_NAMES = {
  'Adagrad':
@@ -75,7 +78,8 @@ def optimize_loss(loss,
  summaries=None,
  colocate_gradients_with_ops=False,
  not_apply_grad_after_first_step=False,
- increment_global_step=True):
+ increment_global_step=True,
+ incr_save=False):
  """Given loss and parameters for optimizer, returns a training op.
 
  Various ways of passing optimizers include:
@@ -146,6 +150,7 @@ def optimize_loss(loss,
  calls `optimize_loss` multiple times per training step (e.g. to optimize
  different parts of the model), use this arg to avoid incrementing
  `global_step` more times than necessary.
+ incr_save: increment dump checkpoints.
 
  Returns:
  Training op.
@@ -300,11 +305,23 @@ def optimize_loss(loss,
 
  # Create gradient updates.
  def _apply_grad():
+ incr_save_ops = []
+ if incr_save:
+ for grad, var in gradients:
+ if isinstance(grad, ops.IndexedSlices):
+ with ops.colocate_with(var):
+ incr_save_op = set_sparse_indices(grad.indices, var_name=var.op.name)
+ incr_save_ops.append(incr_save_op)
+ ops.add_to_collection('SPARSE_UPDATE_VARIABLES', (var, grad.indices.dtype))
+ else:
+ ops.add_to_collection('DENSE_UPDATE_VARIABLES', var)
+
  grad_updates = opt.apply_gradients(
  gradients,
  global_step=global_step if increment_global_step else None,
  name='train')
- return control_flow_ops.with_dependencies([grad_updates], loss)
+
+ return control_flow_ops.with_dependencies([grad_updates] + incr_save_ops, loss)
 
  if not_apply_grad_after_first_step:
  train_tensor = control_flow_ops.cond(global_step > 0, lambda: loss,

easy_rec/python/input/datahub_input.py

@@ -4,7 +4,9 @@
 import time
 
 import numpy as np
+import json
 import tensorflow as tf
+import traceback
 
 from easy_rec.python.input.input import Input
 from easy_rec.python.utils import odps_util
@@ -18,14 +20,18 @@
  from datahub.exceptions import DatahubException
  from datahub.models import RecordType
  from datahub.models import CursorType
-except Exception:
+ import urllib3
+ urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+ logging.getLogger('datahub.account').setLevel(logging.INFO)
+except Exception as ex:
+ # logging.warning(traceback.format_exc(ex))
  logging.warning(
  'DataHub is not installed. You can install it by: pip install pydatahub')
  DataHub = None
 
-
 class DataHubInput(Input):
- """Common IO based interface, could run at local or on data science."""
+ """DataHubInput is used for online train."""
+
 
  def __init__(self,
  data_config,
@@ -35,27 +41,59 @@ def __init__(self,
  task_num=1):
  super(DataHubInput, self).__init__(data_config, feature_config, '',
  task_index, task_num)
- if DataHub is None:
- logging.error('please install datahub: ',
- 'pip install pydatahub ;Python 3.6 recommended')
+
  try:
- self._datahub_config = datahub_config
- if self._datahub_config is None:
- pass
- self._datahub = DataHub(self._datahub_config.akId,
- self._datahub_config.akSecret,
- self._datahub_config.region)
  self._num_epoch = 0
+ self._datahub_config = datahub_config
+ if self._datahub_config is not None:
+ akId = self._datahub_config.akId
+ akSecret = self._datahub_config.akSecret
+ region = self._datahub_config.region
+ if not isinstance(akId, str):
+ akId = akId.encode('utf-8')
+ akSecret = akSecret.encode('utf-8')
+ region = region.encode('utf-8')
+ self._datahub = DataHub(akId, akSecret, region)
+ else:
+ self._datahub = None
  except Exception as ex:
- logging.info('exception in init datahub:', str(ex))
+ logging.info('exception in init datahub: %s' % str(ex))
  pass
+ self._offset_dict = {}
+ if datahub_config:
+ if self._datahub_config.offset_info:
+ self._offset_dict = json.loads(self._datahub_config.offset_info)
+ shard_result = self._datahub.list_shard(self._datahub_config.project,
+ self._datahub_config.topic)
+ shards = shard_result.shards
+ self._shards = [shards[i] for i in range(len(shards)) if (i % task_num) == task_index]
+ logging.info('all shards: %s' % str(self._shards))
+ offset_dict = {}
+ for x in self._shards:
+ if x.shard_id in self._offset_dict:
+ offset_dict[x.shard_id] = self._offset_dict[x.shard_id]
+ self._offset_dict = offset_dict
 
  def _parse_record(self, *fields):
  fields = list(fields)
- inputs = {self._input_fields[x]: fields[x] for x in self._effective_fids}
+ field_dict = {self._input_fields[x]: fields[x] for x in self._effective_fids}
  for x in self._label_fids:
- inputs[self._input_fields[x]] = fields[x]
- return inputs
+ field_dict[self._input_fields[x]] = fields[x]
+ field_dict[Input.DATA_OFFSET] = fields[-1]
+ return field_dict
+
+ def _preprocess(self, field_dict):
+ output_dict = super(DataHubInput, self)._preprocess(field_dict)
+
+ # append offset fields
+ if Input.DATA_OFFSET in field_dict:
+ output_dict[Input.DATA_OFFSET] = field_dict[Input.DATA_OFFSET]
+
+ # for _get_features to include DATA_OFFSET
+ if Input.DATA_OFFSET not in self._appended_fields: 
+ self._appended_fields.append(Input.DATA_OFFSET)
+
+ return output_dict
 
  def _datahub_generator(self):
  logging.info('start epoch[%d]' % self._num_epoch)
@@ -65,62 +103,87 @@ def _datahub_generator(self):
  self.get_type_defaults(x, v)
  for x, v in zip(self._input_field_types, self._input_field_defaults)
  ]
- batch_defaults = [
- np.array([x] * self._data_config.batch_size) for x in record_defaults
+ batch_data = [
+ np.asarray([x] * self._data_config.batch_size, order='C', dtype=object) 
+ if isinstance(x, str) else
+ np.array([x] * self._data_config.batch_size)
+ for x in record_defaults
  ]
+ batch_data.append(json.dumps(self._offset_dict))
+
  try:
  self._datahub.wait_shards_ready(self._datahub_config.project,
  self._datahub_config.topic)
  topic_result = self._datahub.get_topic(self._datahub_config.project,
  self._datahub_config.topic)
  if topic_result.record_type != RecordType.TUPLE:
- logging.error('topic type illegal !')
+ logging.error('datahub topic type(%s) illegal' % str(topic_result.record_type))
  record_schema = topic_result.record_schema
- shard_result = self._datahub.list_shard(self._datahub_config.project,
- self._datahub_config.topic)
- shards = shard_result.shards
- for shard in shards:
- shard_id = shard._shard_id
- cursor_result = self._datahub.get_cursor(self._datahub_config.project,
+
+ batch_size = self._data_config.batch_size
+
+ tid = 0
+ while True:
+ shard_id = self._shards[tid].shard_id
+ tid += 1
+ if tid >= len(self._shards):
+ tid = 0
+ if shard_id not in self._offset_dict:
+ cursor_result = self._datahub.get_cursor(self._datahub_config.project,
  self._datahub_config.topic,
  shard_id, CursorType.OLDEST)
- cursor = cursor_result.cursor
- limit = self._data_config.batch_size
- while True:
- get_result = self._datahub.get_tuple_records(
- self._datahub_config.project, self._datahub_config.topic,
- shard_id, record_schema, cursor, limit)
- batch_data_np = [x.copy() for x in batch_defaults]
- for row_id, record in enumerate(get_result.records):
- for col_id in range(len(record_defaults)):
- if record.values[col_id] not in ['', 'Null', None]:
- batch_data_np[col_id][row_id] = record.values[col_id]
- yield tuple(batch_data_np)
- if 0 == get_result.record_count:
- time.sleep(1)
- cursor = get_result.next_cursor
- except DatahubException as e:
- logging.error(e)
+ cursor = cursor_result.cursor
+ else:
+ cursor = self._offset_dict[shard_id]['cursor']
+
+ get_result = self._datahub.get_tuple_records(
+ self._datahub_config.project, self._datahub_config.topic,
+ shard_id, record_schema, cursor, batch_size)
+ count = get_result.record_count
+ if count == 0:
+ continue
+ time_offset = 0
+ sequence_offset = 0
+ for row_id, record in enumerate(get_result.records):
+ if record.system_time > time_offset:
+ time_offset = record.system_time
+ if record.sequence > sequence_offset:
+ sequence_offset = record.sequence 
+ for col_id in range(len(record_defaults)):
+ if record.values[col_id] not in ['', 'Null', 'null', 'NULL', None]:
+ batch_data[col_id][row_id] = record.values[col_id]
+ else:
+ batch_data[col_id][row_id] = record_defaults[col_id]
+ cursor = get_result.next_cursor
+ self._offset_dict[shard_id] = {'sequence_offset': sequence_offset,
+ 'time_offset': time_offset, 
+ 'cursor': cursor
+ } 
+ batch_data[-1] = json.dumps(self._offset_dict)
+ yield tuple(batch_data)
+ except DatahubException as ex:
+ logging.error('DatahubException: %s' % str(ex))
 
  def _build(self, mode, params):
- # get input type
- list_type = [self.get_tf_type(x) for x in self._input_field_types]
- list_type = tuple(list_type)
- list_shapes = [tf.TensorShape([None]) for x in range(0, len(list_type))]
+ # get input types
+ list_types = [self.get_tf_type(x) for x in self._input_field_types]
+ list_types.append(tf.string)
+ list_types = tuple(list_types)
+ list_shapes = [tf.TensorShape([None]) for x in range(0, len(self._input_field_types))]
+ list_shapes.append(tf.TensorShape([]))
  list_shapes = tuple(list_shapes)
  # read datahub
  dataset = tf.data.Dataset.from_generator(
  self._datahub_generator,
- output_types=list_type,
+ output_types=list_types,
  output_shapes=list_shapes)
  if mode == tf.estimator.ModeKeys.TRAIN:
- dataset = dataset.shuffle(
- self._data_config.shuffle_buffer_size,
- seed=2020,
- reshuffle_each_iteration=True)
- dataset = dataset.repeat(self.num_epochs)
- else:
- dataset = dataset.repeat(1)
+ if self._data_config.shuffle:
+ dataset = dataset.shuffle(
+ self._data_config.shuffle_buffer_size,
+ seed=2020,
+ reshuffle_each_iteration=True)
+
  dataset = dataset.map(
  self._parse_record,
  num_parallel_calls=self._data_config.num_parallel_calls)

easy_rec/python/input/input.py

@@ -23,6 +23,8 @@
 
 class Input(six.with_metaclass(_meta_type, object)):
 
+ DATA_OFFSET = 'DATA_OFFSET' 
+
  def __init__(self,
  data_config,
  feature_configs,