[DIST] Support data sync and evaluation using embedding_lookup_* API

Author: 2sin18

docs/distributed.md

@@ -47,94 +47,59 @@ python -m hybridbackend.run python /path/to/main.py
 
 ## 2. Data Parallelism
 
-HybridBackend provides `hb.functions` and `hb.scope` to rewrite variables and
-optimizers for supporting data paralleism. Also `hb.wraps` could be used for
-user-defined optimizers.
+HybridBackend provides `hb.scope` to rewrite variables and optimizers for
+supporting data paralleism.
 
 ### 2.1 APIs
 
 ```{eval-rst}
-.. autofunction:: hybridbackend.tensorflow.function
 .. autofunction:: hybridbackend.tensorflow.scope
-.. autofunction:: hybridbackend.tensorflow.wraps
 ```
 
-### 2.2 Example: Training inside a function
+### 2.2 Example: Training within a scope
 
 ```python
 import tensorflow as tf
 import hybridbackend.tensorflow as hb
 
-@hb.function(grad_nbuckets=2)
-def foo():
-  # ...
-  loss = tf.losses.get_total_loss()
-  # predefined optimizer
-  opt = tf.train.GradientDescentOptimizer(learning_rate=lr)
-```
-
-### 2.3 Example: Training within a scope
-
-```python
-import tensorflow as tf
-import hybridbackend.tensorflow as hb
-
-with hb.scope(grad_nbuckets=2):
+with hb.scope():
   # ...
   loss = tf.losses.get_total_loss()
   # predefined optimizer
   opt = tf.train.GradientDescentOptimizer(learning_rate=lr)
 ```
 
-### 2.4 Example: Defines a new optimizer
-
-```python
-import tensorflow as tf
-import hybridbackend.tensorflow as hb
-
-@hb.wraps
-class MyOptimizer(tf.train.Optimizer):
-  # ...
-
-# ...
-def foo():
-  loss = tf.losses.get_total_loss()
-  opt = MyOptimizer(learning_rate=lr)
-  train_op = opt.minimize(loss)
-```
-
-## 2. Model Parallelism
+## 2. Embedding-Sharded Data Parallelism
 
 HybridBackend provides option `sharding` to shard variables and support
-model-parallel embedding layers.
+embedding-sharded data paralleism.
 
 ### 2.1 APIs
 
 ```{eval-rst}
-.. autofunction:: hybridbackend.tensorflow.data.make_one_shot_iterator
-.. autofunction:: hybridbackend.tensorflow.data.make_initializable_iterator
-.. autofunction:: hybridbackend.tensorflow.train.export
 .. autofunction:: hybridbackend.tensorflow.metrics.accuracy
 .. autofunction:: hybridbackend.tensorflow.metrics.auc
+.. autofunction:: hybridbackend.tensorflow.train.EvaluationHook
+.. autofunction:: hybridbackend.tensorflow.train.export
 ```
 
-### 2.2 Example: Sharding inside a function
+### 2.2 Example: Sharding embedding weights within a scope
 
 ```python
 import tensorflow as tf
 import hybridbackend.tensorflow as hb
 
-@hb.function(grad_nbuckets=2)
 def foo():
   # ...
-  with hb.scope(sharding=True):
-    embedding_weights = tf.get_variable(
-      'emb_weights', shape=[bucket_size, dim_size])
-  embedding = tf.nn.embedding_lookup(embedding_weights, ids)
-  # ...
-  loss = tf.losses.get_total_loss()
-  # predefined optimizer
-  opt = tf.train.GradientDescentOptimizer(learning_rate=lr)
+  with hb.scope():
+    with hb.scope(sharding=True):
+      embedding_weights = tf.get_variable(
+        'emb_weights', shape=[bucket_size, dim_size])
+    embedding = tf.nn.embedding_lookup(embedding_weights, ids)
+    # ...
+    loss = tf.losses.get_total_loss()
+    # predefined optimizer
+    opt = tf.train.GradientDescentOptimizer(learning_rate=lr)
 ```
 
 ### 2.3 Example: Evaluation
@@ -152,19 +117,16 @@ def eval_fn():
 
 with tf.Graph().as_default():
   with hb.scope():
-    batch = hb.data.make_one_shot_iterator(train_ds).get_next()
+    batch = tf.data.make_one_shot_iterator(train_ds).get_next()
     # ...
     with hb.scope(sharding=True):
       embedding_weights = tf.get_variable(
         'emb_weights', shape=[bucket_size, dim_size])
     embedding = tf.nn.embedding_lookup(embedding_weights, ids)
     # ...
+    hooks.append(hb.train.EvaluationHook(eval_fn, every_n_iter=1000))
 
-    with tf.train.MonitoredTrainingSession(
-        '',
-        hooks=hooks,
-        eval_every_n_iter=1000,
-        eval_fn=eval_fn) as sess:
+    with tf.train.MonitoredTrainingSession('', hooks=hooks) as sess:
       while not sess.should_stop():
         sess.run(train_op)
 ```

docs/high_level_api.md

@@ -134,29 +134,3 @@ def model_fn(features, labels, mode, params):
       loss=loss,
       train_op=train_op)
 ```
-
-### 2.4 Example: Use user-defined embedding backend
-
-```python
-import tensorflow as tf
-import hybridbackend.tensorflow as hb
-
-class MyEmbeddingBackend(hb.embedding.Backend):
-  NAME = 'MYEMB'
-
-  def build(self, column, name, shape, **kwargs):
-    return mymodule.get_my_emb(name, shape, **kwargs)
-
-  def lookup(self, column, weight, inputs, sharded=False, buffered=False):
-    r'''Lookup for embedding vectors.
-    '''
-    return mymodule.lookup(weight, inputs)
-
-hb.embedding.Backend.register(MyEmbeddingBackend())
-```
-
-```python
-from mymodule import MyEmbeddingBackend
-
-hb.context.options.emb_backend = 'PAIEV'
-```

docs/images

@@ -7,5 +7,4 @@ introduction
 architecture
 data
 distributed
-high_level_api
 ```

images/architecture.png docs/images/architecture.pngimages/architecture.png renamed to docs/images/architecture.png

@@ -123,7 +123,7 @@
     "data_spec = DataSpec.read('ranking/taobao/data/spec.json')\n",
     "\n",
     "\n",
-    "def train(iterator, embedding_weight_device, dnn_device, hooks):\n",
+    "def train(iterator, embedding_weight_device, dnn_device):\n",
     "  batch = iterator.get_next()\n",
     "  batch.pop('ts')\n",
     "  labels = tf.reshape(tf.to_float(batch.pop('label')), shape=[-1, 1])\n",
@@ -150,6 +150,7 @@
     "    opt = tf.train.AdagradOptimizer(learning_rate=0.001)\n",
     "    train_op = opt.minimize(loss, global_step=step)\n",
     "\n",
+    "  hooks = []\n",
     "  hooks.append(tf.train.StepCounterHook(10))\n",
     "  hooks.append(tf.train.StopAtStepHook(train_max_steps))\n",
     "  config = tf.ConfigProto(allow_soft_placement=True)\n",
@@ -203,7 +204,7 @@
     "  iterator = tf.data.make_one_shot_iterator(ds)\n",
     "\n",
     "  with tf.device('/gpu:0'):\n",
-    "    train(iterator, '/cpu:0', '/gpu:0', [])"
+    "    train(iterator, '/cpu:0', '/gpu:0')"
    ]
   },
   {
@@ -239,7 +240,7 @@
     "  iterator = tf.data.make_one_shot_iterator(ds)\n",
     "\n",
     "  with tf.device('/gpu:0'):\n",
-    "    train(iterator, '/cpu:0', '/gpu:0', [])"
+    "    train(iterator, '/cpu:0', '/gpu:0')"
    ]
   },
   {
@@ -290,8 +291,7 @@
     "  iterator = tf.data.make_one_shot_iterator(ds)\n",
     "\n",
     "  with tf.device('/gpu:0'):\n",
-    "    iterator = hb.data.Iterator(iterator, 2)\n",
-    "    train(iterator, '/cpu:0', '/gpu:0', [hb.data.Iterator.Hook()])"
+    "    train(iterator, '/cpu:0', '/gpu:0')"
    ]
   }
  ],