We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Simplify user API: don't rely on column name to find labels.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `edl_k8s_examples/mnist_model.py`
Content:
```
1 import tensorflow as tf
2 import numpy as np
3 
4 
5 class MnistModel(tf.keras.Model):
6     def __init__(self, channel_last=True):
7         super(MnistModel, self).__init__(name='mnist_model')
8         if channel_last:
9             self._reshape = tf.keras.layers.Reshape((28, 28, 1))
10         else:
11             self._reshape = tf.keras.layers.Reshape((1, 28, 28))
12         self._conv1 = tf.keras.layers.Conv2D(
13             32, kernel_size=(3, 3), activation='relu')
14         self._conv2 = tf.keras.layers.Conv2D(
15             64, kernel_size=(3, 3), activation='relu')
16         self._batch_norm = tf.keras.layers.BatchNormalization()
17         self._maxpooling = tf.keras.layers.MaxPooling2D(
18             pool_size=(2, 2))
19         self._dropout = tf.keras.layers.Dropout(0.25)
20         self._flatten = tf.keras.layers.Flatten()
21         self._dense = tf.keras.layers.Dense(10)
22 
23     def call(self, inputs, training=False):
24         x = self._reshape(inputs)
25         x = self._conv1(x)
26         x = self._conv2(x)
27         x = self._batch_norm(x, training=training)
28         x = self._maxpooling(x)
29         if training:
30             x = self._dropout(x, training=training)
31         x = self._flatten(x)
32         x = self._dense(x)
33         return x
34 
35 
36 model = MnistModel()
37 model.build(input_shape=(1, 28, 28))
38 
39 input_names = ['image']
40         
41 def loss(output, labels):
42     return tf.reduce_mean(
43         tf.nn.sparse_softmax_cross_entropy_with_logits(
44             logits=output, labels=labels['label']))
45 
46 def optimizer(lr=0.1):
47     return tf.train.GradientDescentOptimizer(lr)
48 
49 def input_fn(records):
50     image_list = []
51     label_list = []
52     # deserialize
53     for r in records:
54         parsed = np.frombuffer(r, dtype="uint8")
55         label = parsed[-1]
56         image = np.resize(parsed[:-1], new_shape=(28, 28))
57         image = image.astype(np.float32)
58         image /= 255
59         label = label.astype(np.int32)
60         image_list.append(image)
61         label_list.append(label)
62 
63     # batching
64     batch_size = len(image_list)
65     images = np.concatenate(image_list, axis=0)
66     images = np.reshape(images, (batch_size, 28, 28))
67     labels = np.array(label_list)
68     return {'image': images, 'label': labels}
69
```

Path: `elasticdl/worker/worker.py`
Content:
```
1 import tensorflow as tf
2 assert tf.executing_eagerly()
3 
4 from google.protobuf import empty_pb2
5 from tensorflow.python.ops import math_ops
6 from elasticdl.proto import master_pb2_grpc
7 from elasticdl.proto import master_pb2
8 from elasticdl.common.ndarray import ndarray_to_tensor, tensor_to_ndarray
9 from elasticdl.common.model_helper import load_user_model
10 import itertools
11 import recordio
12 
13 # the default max number of a minibatch retrain as its gradients are not accepted by master.
14 DEFAULT_MAX_MINIBATCH_RETRAIN_NUM = 64
15 
16 class Worker(object):
17     """ElasticDL worker"""
18 
19     def __init__(self,
20                  model_file,
21                  channel=None,
22                  max_retrain_num=DEFAULT_MAX_MINIBATCH_RETRAIN_NUM):
23         """
24         Arguments:
25             model_module: A module to define the model
26             channel: grpc channel
27             max_retrain_num: max number of a minibatch retrain as its gradients are not accepted by master
28         """
29 
30         model_module = load_user_model(model_file)
31         self._model = model_module.model
32         self._input_fn = model_module.input_fn 
33         self._opt_fn = model_module.optimizer
34         self._loss = model_module.loss
35         self._input_names = model_module.input_names
36 
37         if channel is None:
38             self._stub = None
39         else:
40             self._stub = master_pb2_grpc.MasterStub(channel)
41         self._max_retrain_num = max_retrain_num
42         self._model_version = -1
43 
44     def get_task(self):
45         """
46         get task from master
47         """
48         return self._stub.GetTask(empty_pb2.Empty())
49 
50     def get_model(self, min_version):
51         """
52         get model from master, and update model_version
53         """
54         req = master_pb2.GetModelRequest()
55         req.min_version = min_version
56         model = self._stub.GetModel(req)
57 
58         for var in self._model.trainable_variables:
59             # Assumes all trainable variables exist in model.param.
60             var.assign(
61                 tensor_to_ndarray(model.param[var.name]))
62         self._model_version = model.version
63 
64     def report_task_result(self, task_id, err_msg):
65         """
66         report task result to master
67         """
68         report = master_pb2.ReportTaskResultRequest()
69         report.task_id = task_id
70         report.err_message = err_msg
71         return self._stub.ReportTaskResult(report)
72 
73     def report_gradient(self, grads):
74         """
75         report gradient to ps, return (accepted, model_version) from rpc call.
76         """
77         req = master_pb2.ReportGradientRequest()
78         for g, v in zip(grads, self._model.trainable_variables):
79             req.gradient[v.name].CopyFrom(
80                 ndarray_to_tensor(g.numpy()))
81         req.model_version = self._model_version
82         res = self._stub.ReportGradient(req)
83         return res.accepted, res.model_version
84 
85     def distributed_train(self):
86         """
87         Distributed training.
88         """
89         while True:
90             task = self.get_task()
91             if not task.shard_file_name:
92                 # No more task
93                 break
94             batch_size = task.minibatch_size
95             err_msg = ""
96             try:
97                 with recordio.File(task.shard_file_name, "r") as rdio_r:
98                     reader = rdio_r.get_reader(task.start, task.end)
99                     min_model_version = task.model_version
100                     while True:
101                         record_buf = list(
102                             itertools.islice(reader, 0, batch_size))
103                         if not record_buf:
104                             break
105 
106                         for _ in range(self._max_retrain_num):
107                             # TODO: optimize the logic to avoid unnecessary get_model call.
108                             self.get_model(
109                                 max(self._model_version, min_model_version))
110 
111                             batch_input_data = self._input_fn(record_buf)
112 
113                             with tf.GradientTape() as tape:
114                                 inputs = []
115                                 for input_name in self._input_names:
116                                     inputs.append(batch_input_data[input_name])
117                                 if len(inputs) == 1:
118                                     inputs = inputs[0]
119                                 outputs = self._model.call(inputs, training=True)
120                                 loss = self._loss(outputs, batch_input_data)
121 
122                                 # TODO:  Add regularization loss if any,
123                                 #        which should be divided by the number of contributing workers.
124                             grads = tape.gradient(
125                                 loss, self._model.trainable_variables)
126                             print("Loss is ", loss.numpy())
127 
128                             accepted, min_model_version = self.report_gradient(
129                                 grads)
130                             if accepted:
131                                 break
132                         else:
133                             # Worker got stuck, fail the task.
134                             # TODO: stop the worker if it fails to make any progress for some time.
135                             raise RuntimeError("Worker got stuck")
136 
137 
138             except Exception as ex:
139                 err_msg = str(ex)
140             self.report_task_result(task.task_id, err_msg)
141 
142     def local_train(self, file_list, batch_size, epoch=1, kwargs=None):
143         """
144         Local training for local testing. Must in eager mode.
145         Argments:
146             batch_size: batch size in training
147             epoch: the number of epoch in training
148             kwargs: contains a dict of parameters used in training
149         """
150         optimizer = self._opt_fn()
151         for _ in range(epoch):
152             for f in file_list:
153                 with recordio.File(f, "r") as rdio_r:
154                     reader = rdio_r.get_reader(0, rdio_r.count())
155                     while True:
156                         record_buf = list(
157                             itertools.islice(reader, 0, batch_size))
158                         if not record_buf:
159                             break
160 
161                         data = self._input_fn(record_buf)
162 
163                         with tf.GradientTape() as tape:
164                             inputs = []
165                             for input_name in self._input_names:
166                                 inputs.append(data[input_name])
167                             if len(inputs) == 1:
168                                 inputs = inputs[0]
169                             outputs = self._model.call(inputs, training=True)
170                             loss = self._loss(outputs, data)
171 
172                             # Add regularization loss if any.
173                             # Note: for distributed training, the regularization loss should
174                             #       be divided by the number of contributing workers, which
175                             #       might be difficult for elasticdl.
176                             if self._model.losses:
177                                 loss += math_ops.add_n(self._model.losses)
178                         grads = tape.gradient(
179                             loss, self._model.trainable_variables)
180                         optimizer.apply_gradients(
181                             zip(grads, self._model.trainable_variables))
182                         print("Loss is ", loss.numpy())
183
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/edl_k8s_examples/mnist_model.py b/edl_k8s_examples/mnist_model.py
--- a/edl_k8s_examples/mnist_model.py
+++ b/edl_k8s_examples/mnist_model.py
@@ -41,7 +41,7 @@
 def loss(output, labels):
     return tf.reduce_mean(
         tf.nn.sparse_softmax_cross_entropy_with_logits(
-            logits=output, labels=labels['label']))
+            logits=output, labels=labels))
 
 def optimizer(lr=0.1):
     return tf.train.GradientDescentOptimizer(lr)
@@ -65,4 +65,4 @@
     images = np.concatenate(image_list, axis=0)
     images = np.reshape(images, (batch_size, 28, 28))
     labels = np.array(label_list)
-    return {'image': images, 'label': labels}
+    return ({'image': images}, labels)
diff --git a/elasticdl/worker/worker.py b/elasticdl/worker/worker.py
--- a/elasticdl/worker/worker.py
+++ b/elasticdl/worker/worker.py
@@ -108,7 +108,7 @@
                             self.get_model(
                                 max(self._model_version, min_model_version))
 
-                            batch_input_data = self._input_fn(record_buf)
+                            batch_input_data, batch_label = self._input_fn(record_buf)
 
                             with tf.GradientTape() as tape:
                                 inputs = []
@@ -117,7 +117,7 @@
                                 if len(inputs) == 1:
                                     inputs = inputs[0]
                                 outputs = self._model.call(inputs, training=True)
-                                loss = self._loss(outputs, batch_input_data)
+                                loss = self._loss(outputs, batch_label)
 
                                 # TODO:  Add regularization loss if any,
                                 #        which should be divided by the number of contributing workers.
@@ -158,7 +158,7 @@
                         if not record_buf:
                             break
 
-                        data = self._input_fn(record_buf)
+                        data, labels = self._input_fn(record_buf)
 
                         with tf.GradientTape() as tape:
                             inputs = []
@@ -167,7 +167,7 @@
                             if len(inputs) == 1:
                                 inputs = inputs[0]
                             outputs = self._model.call(inputs, training=True)
-                            loss = self._loss(outputs, data)
+                            loss = self._loss(outputs, labels)
 
                             # Add regularization loss if any.
                             # Note: for distributed training, the regularization loss should

{"golden_diff": "diff --git a/edl_k8s_examples/mnist_model.py b/edl_k8s_examples/mnist_model.py\n--- a/edl_k8s_examples/mnist_model.py\n+++ b/edl_k8s_examples/mnist_model.py\n@@ -41,7 +41,7 @@\n def loss(output, labels):\n     return tf.reduce_mean(\n         tf.nn.sparse_softmax_cross_entropy_with_logits(\n-            logits=output, labels=labels['label']))\n+            logits=output, labels=labels))\n \n def optimizer(lr=0.1):\n     return tf.train.GradientDescentOptimizer(lr)\n@@ -65,4 +65,4 @@\n     images = np.concatenate(image_list, axis=0)\n     images = np.reshape(images, (batch_size, 28, 28))\n     labels = np.array(label_list)\n-    return {'image': images, 'label': labels}\n+    return ({'image': images}, labels)\ndiff --git a/elasticdl/worker/worker.py b/elasticdl/worker/worker.py\n--- a/elasticdl/worker/worker.py\n+++ b/elasticdl/worker/worker.py\n@@ -108,7 +108,7 @@\n                             self.get_model(\n                                 max(self._model_version, min_model_version))\n \n-                            batch_input_data = self._input_fn(record_buf)\n+                            batch_input_data, batch_label = self._input_fn(record_buf)\n \n                             with tf.GradientTape() as tape:\n                                 inputs = []\n@@ -117,7 +117,7 @@\n                                 if len(inputs) == 1:\n                                     inputs = inputs[0]\n                                 outputs = self._model.call(inputs, training=True)\n-                                loss = self._loss(outputs, batch_input_data)\n+                                loss = self._loss(outputs, batch_label)\n \n                                 # TODO:  Add regularization loss if any,\n                                 #        which should be divided by the number of contributing workers.\n@@ -158,7 +158,7 @@\n                         if not record_buf:\n                             break\n \n-                        data = self._input_fn(record_buf)\n+                        data, labels = self._input_fn(record_buf)\n \n                         with tf.GradientTape() as tape:\n                             inputs = []\n@@ -167,7 +167,7 @@\n                             if len(inputs) == 1:\n                                 inputs = inputs[0]\n                             outputs = self._model.call(inputs, training=True)\n-                            loss = self._loss(outputs, data)\n+                            loss = self._loss(outputs, labels)\n \n                             # Add regularization loss if any.\n                             # Note: for distributed training, the regularization loss should\n", "issue": "Simplify user API: don't rely on column name to find labels.\n\n", "before_files": [{"content": "import tensorflow as tf\nimport numpy as np\n\n\nclass MnistModel(tf.keras.Model):\n    def __init__(self, channel_last=True):\n        super(MnistModel, self).__init__(name='mnist_model')\n        if channel_last:\n            self._reshape = tf.keras.layers.Reshape((28, 28, 1))\n        else:\n            self._reshape = tf.keras.layers.Reshape((1, 28, 28))\n        self._conv1 = tf.keras.layers.Conv2D(\n            32, kernel_size=(3, 3), activation='relu')\n        self._conv2 = tf.keras.layers.Conv2D(\n            64, kernel_size=(3, 3), activation='relu')\n        self._batch_norm = tf.keras.layers.BatchNormalization()\n        self._maxpooling = tf.keras.layers.MaxPooling2D(\n            pool_size=(2, 2))\n        self._dropout = tf.keras.layers.Dropout(0.25)\n        self._flatten = tf.keras.layers.Flatten()\n        self._dense = tf.keras.layers.Dense(10)\n\n    def call(self, inputs, training=False):\n        x = self._reshape(inputs)\n        x = self._conv1(x)\n        x = self._conv2(x)\n        x = self._batch_norm(x, training=training)\n        x = self._maxpooling(x)\n        if training:\n            x = self._dropout(x, training=training)\n        x = self._flatten(x)\n        x = self._dense(x)\n        return x\n\n\nmodel = MnistModel()\nmodel.build(input_shape=(1, 28, 28))\n\ninput_names = ['image']\n        \ndef loss(output, labels):\n    return tf.reduce_mean(\n        tf.nn.sparse_softmax_cross_entropy_with_logits(\n            logits=output, labels=labels['label']))\n\ndef optimizer(lr=0.1):\n    return tf.train.GradientDescentOptimizer(lr)\n\ndef input_fn(records):\n    image_list = []\n    label_list = []\n    # deserialize\n    for r in records:\n        parsed = np.frombuffer(r, dtype=\"uint8\")\n        label = parsed[-1]\n        image = np.resize(parsed[:-1], new_shape=(28, 28))\n        image = image.astype(np.float32)\n        image /= 255\n        label = label.astype(np.int32)\n        image_list.append(image)\n        label_list.append(label)\n\n    # batching\n    batch_size = len(image_list)\n    images = np.concatenate(image_list, axis=0)\n    images = np.reshape(images, (batch_size, 28, 28))\n    labels = np.array(label_list)\n    return {'image': images, 'label': labels}\n", "path": "edl_k8s_examples/mnist_model.py"}, {"content": "import tensorflow as tf\nassert tf.executing_eagerly()\n\nfrom google.protobuf import empty_pb2\nfrom tensorflow.python.ops import math_ops\nfrom elasticdl.proto import master_pb2_grpc\nfrom elasticdl.proto import master_pb2\nfrom elasticdl.common.ndarray import ndarray_to_tensor, tensor_to_ndarray\nfrom elasticdl.common.model_helper import load_user_model\nimport itertools\nimport recordio\n\n# the default max number of a minibatch retrain as its gradients are not accepted by master.\nDEFAULT_MAX_MINIBATCH_RETRAIN_NUM = 64\n\nclass Worker(object):\n    \"\"\"ElasticDL worker\"\"\"\n\n    def __init__(self,\n                 model_file,\n                 channel=None,\n                 max_retrain_num=DEFAULT_MAX_MINIBATCH_RETRAIN_NUM):\n        \"\"\"\n        Arguments:\n            model_module: A module to define the model\n            channel: grpc channel\n            max_retrain_num: max number of a minibatch retrain as its gradients are not accepted by master\n        \"\"\"\n\n        model_module = load_user_model(model_file)\n        self._model = model_module.model\n        self._input_fn = model_module.input_fn \n        self._opt_fn = model_module.optimizer\n        self._loss = model_module.loss\n        self._input_names = model_module.input_names\n\n        if channel is None:\n            self._stub = None\n        else:\n            self._stub = master_pb2_grpc.MasterStub(channel)\n        self._max_retrain_num = max_retrain_num\n        self._model_version = -1\n\n    def get_task(self):\n        \"\"\"\n        get task from master\n        \"\"\"\n        return self._stub.GetTask(empty_pb2.Empty())\n\n    def get_model(self, min_version):\n        \"\"\"\n        get model from master, and update model_version\n        \"\"\"\n        req = master_pb2.GetModelRequest()\n        req.min_version = min_version\n        model = self._stub.GetModel(req)\n\n        for var in self._model.trainable_variables:\n            # Assumes all trainable variables exist in model.param.\n            var.assign(\n                tensor_to_ndarray(model.param[var.name]))\n        self._model_version = model.version\n\n    def report_task_result(self, task_id, err_msg):\n        \"\"\"\n        report task result to master\n        \"\"\"\n        report = master_pb2.ReportTaskResultRequest()\n        report.task_id = task_id\n        report.err_message = err_msg\n        return self._stub.ReportTaskResult(report)\n\n    def report_gradient(self, grads):\n        \"\"\"\n        report gradient to ps, return (accepted, model_version) from rpc call.\n        \"\"\"\n        req = master_pb2.ReportGradientRequest()\n        for g, v in zip(grads, self._model.trainable_variables):\n            req.gradient[v.name].CopyFrom(\n                ndarray_to_tensor(g.numpy()))\n        req.model_version = self._model_version\n        res = self._stub.ReportGradient(req)\n        return res.accepted, res.model_version\n\n    def distributed_train(self):\n        \"\"\"\n        Distributed training.\n        \"\"\"\n        while True:\n            task = self.get_task()\n            if not task.shard_file_name:\n                # No more task\n                break\n            batch_size = task.minibatch_size\n            err_msg = \"\"\n            try:\n                with recordio.File(task.shard_file_name, \"r\") as rdio_r:\n                    reader = rdio_r.get_reader(task.start, task.end)\n                    min_model_version = task.model_version\n                    while True:\n                        record_buf = list(\n                            itertools.islice(reader, 0, batch_size))\n                        if not record_buf:\n                            break\n\n                        for _ in range(self._max_retrain_num):\n                            # TODO: optimize the logic to avoid unnecessary get_model call.\n                            self.get_model(\n                                max(self._model_version, min_model_version))\n\n                            batch_input_data = self._input_fn(record_buf)\n\n                            with tf.GradientTape() as tape:\n                                inputs = []\n                                for input_name in self._input_names:\n                                    inputs.append(batch_input_data[input_name])\n                                if len(inputs) == 1:\n                                    inputs = inputs[0]\n                                outputs = self._model.call(inputs, training=True)\n                                loss = self._loss(outputs, batch_input_data)\n\n                                # TODO:  Add regularization loss if any,\n                                #        which should be divided by the number of contributing workers.\n                            grads = tape.gradient(\n                                loss, self._model.trainable_variables)\n                            print(\"Loss is \", loss.numpy())\n\n                            accepted, min_model_version = self.report_gradient(\n                                grads)\n                            if accepted:\n                                break\n                        else:\n                            # Worker got stuck, fail the task.\n                            # TODO: stop the worker if it fails to make any progress for some time.\n                            raise RuntimeError(\"Worker got stuck\")\n\n\n            except Exception as ex:\n                err_msg = str(ex)\n            self.report_task_result(task.task_id, err_msg)\n\n    def local_train(self, file_list, batch_size, epoch=1, kwargs=None):\n        \"\"\"\n        Local training for local testing. Must in eager mode.\n        Argments:\n            batch_size: batch size in training\n            epoch: the number of epoch in training\n            kwargs: contains a dict of parameters used in training\n        \"\"\"\n        optimizer = self._opt_fn()\n        for _ in range(epoch):\n            for f in file_list:\n                with recordio.File(f, \"r\") as rdio_r:\n                    reader = rdio_r.get_reader(0, rdio_r.count())\n                    while True:\n                        record_buf = list(\n                            itertools.islice(reader, 0, batch_size))\n                        if not record_buf:\n                            break\n\n                        data = self._input_fn(record_buf)\n\n                        with tf.GradientTape() as tape:\n                            inputs = []\n                            for input_name in self._input_names:\n                                inputs.append(data[input_name])\n                            if len(inputs) == 1:\n                                inputs = inputs[0]\n                            outputs = self._model.call(inputs, training=True)\n                            loss = self._loss(outputs, data)\n\n                            # Add regularization loss if any.\n                            # Note: for distributed training, the regularization loss should\n                            #       be divided by the number of contributing workers, which\n                            #       might be difficult for elasticdl.\n                            if self._model.losses:\n                                loss += math_ops.add_n(self._model.losses)\n                        grads = tape.gradient(\n                            loss, self._model.trainable_variables)\n                        optimizer.apply_gradients(\n                            zip(grads, self._model.trainable_variables))\n                        print(\"Loss is \", loss.numpy())\n", "path": "elasticdl/worker/worker.py"}], "after_files": [{"content": "import tensorflow as tf\nimport numpy as np\n\n\nclass MnistModel(tf.keras.Model):\n    def __init__(self, channel_last=True):\n        super(MnistModel, self).__init__(name='mnist_model')\n        if channel_last:\n            self._reshape = tf.keras.layers.Reshape((28, 28, 1))\n        else:\n            self._reshape = tf.keras.layers.Reshape((1, 28, 28))\n        self._conv1 = tf.keras.layers.Conv2D(\n            32, kernel_size=(3, 3), activation='relu')\n        self._conv2 = tf.keras.layers.Conv2D(\n            64, kernel_size=(3, 3), activation='relu')\n        self._batch_norm = tf.keras.layers.BatchNormalization()\n        self._maxpooling = tf.keras.layers.MaxPooling2D(\n            pool_size=(2, 2))\n        self._dropout = tf.keras.layers.Dropout(0.25)\n        self._flatten = tf.keras.layers.Flatten()\n        self._dense = tf.keras.layers.Dense(10)\n\n    def call(self, inputs, training=False):\n        x = self._reshape(inputs)\n        x = self._conv1(x)\n        x = self._conv2(x)\n        x = self._batch_norm(x, training=training)\n        x = self._maxpooling(x)\n        if training:\n            x = self._dropout(x, training=training)\n        x = self._flatten(x)\n        x = self._dense(x)\n        return x\n\n\nmodel = MnistModel()\nmodel.build(input_shape=(1, 28, 28))\n\ninput_names = ['image']\n        \ndef loss(output, labels):\n    return tf.reduce_mean(\n        tf.nn.sparse_softmax_cross_entropy_with_logits(\n            logits=output, labels=labels))\n\ndef optimizer(lr=0.1):\n    return tf.train.GradientDescentOptimizer(lr)\n\ndef input_fn(records):\n    image_list = []\n    label_list = []\n    # deserialize\n    for r in records:\n        parsed = np.frombuffer(r, dtype=\"uint8\")\n        label = parsed[-1]\n        image = np.resize(parsed[:-1], new_shape=(28, 28))\n        image = image.astype(np.float32)\n        image /= 255\n        label = label.astype(np.int32)\n        image_list.append(image)\n        label_list.append(label)\n\n    # batching\n    batch_size = len(image_list)\n    images = np.concatenate(image_list, axis=0)\n    images = np.reshape(images, (batch_size, 28, 28))\n    labels = np.array(label_list)\n    return ({'image': images}, labels)\n", "path": "edl_k8s_examples/mnist_model.py"}, {"content": "import tensorflow as tf\nassert tf.executing_eagerly()\n\nfrom google.protobuf import empty_pb2\nfrom tensorflow.python.ops import math_ops\nfrom elasticdl.proto import master_pb2_grpc\nfrom elasticdl.proto import master_pb2\nfrom elasticdl.common.ndarray import ndarray_to_tensor, tensor_to_ndarray\nfrom elasticdl.common.model_helper import load_user_model\nimport itertools\nimport recordio\n\n# the default max number of a minibatch retrain as its gradients are not accepted by master.\nDEFAULT_MAX_MINIBATCH_RETRAIN_NUM = 64\n\nclass Worker(object):\n    \"\"\"ElasticDL worker\"\"\"\n\n    def __init__(self,\n                 model_file,\n                 channel=None,\n                 max_retrain_num=DEFAULT_MAX_MINIBATCH_RETRAIN_NUM):\n        \"\"\"\n        Arguments:\n            model_module: A module to define the model\n            channel: grpc channel\n            max_retrain_num: max number of a minibatch retrain as its gradients are not accepted by master\n        \"\"\"\n\n        model_module = load_user_model(model_file)\n        self._model = model_module.model\n        self._input_fn = model_module.input_fn \n        self._opt_fn = model_module.optimizer\n        self._loss = model_module.loss\n        self._input_names = model_module.input_names\n\n        if channel is None:\n            self._stub = None\n        else:\n            self._stub = master_pb2_grpc.MasterStub(channel)\n        self._max_retrain_num = max_retrain_num\n        self._model_version = -1\n\n    def get_task(self):\n        \"\"\"\n        get task from master\n        \"\"\"\n        return self._stub.GetTask(empty_pb2.Empty())\n\n    def get_model(self, min_version):\n        \"\"\"\n        get model from master, and update model_version\n        \"\"\"\n        req = master_pb2.GetModelRequest()\n        req.min_version = min_version\n        model = self._stub.GetModel(req)\n\n        for var in self._model.trainable_variables:\n            # Assumes all trainable variables exist in model.param.\n            var.assign(\n                tensor_to_ndarray(model.param[var.name]))\n        self._model_version = model.version\n\n    def report_task_result(self, task_id, err_msg):\n        \"\"\"\n        report task result to master\n        \"\"\"\n        report = master_pb2.ReportTaskResultRequest()\n        report.task_id = task_id\n        report.err_message = err_msg\n        return self._stub.ReportTaskResult(report)\n\n    def report_gradient(self, grads):\n        \"\"\"\n        report gradient to ps, return (accepted, model_version) from rpc call.\n        \"\"\"\n        req = master_pb2.ReportGradientRequest()\n        for g, v in zip(grads, self._model.trainable_variables):\n            req.gradient[v.name].CopyFrom(\n                ndarray_to_tensor(g.numpy()))\n        req.model_version = self._model_version\n        res = self._stub.ReportGradient(req)\n        return res.accepted, res.model_version\n\n    def distributed_train(self):\n        \"\"\"\n        Distributed training.\n        \"\"\"\n        while True:\n            task = self.get_task()\n            if not task.shard_file_name:\n                # No more task\n                break\n            batch_size = task.minibatch_size\n            err_msg = \"\"\n            try:\n                with recordio.File(task.shard_file_name, \"r\") as rdio_r:\n                    reader = rdio_r.get_reader(task.start, task.end)\n                    min_model_version = task.model_version\n                    while True:\n                        record_buf = list(\n                            itertools.islice(reader, 0, batch_size))\n                        if not record_buf:\n                            break\n\n                        for _ in range(self._max_retrain_num):\n                            # TODO: optimize the logic to avoid unnecessary get_model call.\n                            self.get_model(\n                                max(self._model_version, min_model_version))\n\n                            batch_input_data, batch_label = self._input_fn(record_buf)\n\n                            with tf.GradientTape() as tape:\n                                inputs = []\n                                for input_name in self._input_names:\n                                    inputs.append(batch_input_data[input_name])\n                                if len(inputs) == 1:\n                                    inputs = inputs[0]\n                                outputs = self._model.call(inputs, training=True)\n                                loss = self._loss(outputs, batch_label)\n\n                                # TODO:  Add regularization loss if any,\n                                #        which should be divided by the number of contributing workers.\n                            grads = tape.gradient(\n                                loss, self._model.trainable_variables)\n                            print(\"Loss is \", loss.numpy())\n\n                            accepted, min_model_version = self.report_gradient(\n                                grads)\n                            if accepted:\n                                break\n                        else:\n                            # Worker got stuck, fail the task.\n                            # TODO: stop the worker if it fails to make any progress for some time.\n                            raise RuntimeError(\"Worker got stuck\")\n\n\n            except Exception as ex:\n                err_msg = str(ex)\n            self.report_task_result(task.task_id, err_msg)\n\n    def local_train(self, file_list, batch_size, epoch=1, kwargs=None):\n        \"\"\"\n        Local training for local testing. Must in eager mode.\n        Argments:\n            batch_size: batch size in training\n            epoch: the number of epoch in training\n            kwargs: contains a dict of parameters used in training\n        \"\"\"\n        optimizer = self._opt_fn()\n        for _ in range(epoch):\n            for f in file_list:\n                with recordio.File(f, \"r\") as rdio_r:\n                    reader = rdio_r.get_reader(0, rdio_r.count())\n                    while True:\n                        record_buf = list(\n                            itertools.islice(reader, 0, batch_size))\n                        if not record_buf:\n                            break\n\n                        data, labels = self._input_fn(record_buf)\n\n                        with tf.GradientTape() as tape:\n                            inputs = []\n                            for input_name in self._input_names:\n                                inputs.append(data[input_name])\n                            if len(inputs) == 1:\n                                inputs = inputs[0]\n                            outputs = self._model.call(inputs, training=True)\n                            loss = self._loss(outputs, labels)\n\n                            # Add regularization loss if any.\n                            # Note: for distributed training, the regularization loss should\n                            #       be divided by the number of contributing workers, which\n                            #       might be difficult for elasticdl.\n                            if self._model.losses:\n                                loss += math_ops.add_n(self._model.losses)\n                        grads = tape.gradient(\n                            loss, self._model.trainable_variables)\n                        optimizer.apply_gradients(\n                            zip(grads, self._model.trainable_variables))\n                        print(\"Loss is \", loss.numpy())\n", "path": "elasticdl/worker/worker.py"}]}