We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
non-trainable variables update in training
Some keras layers have non-trainable variables, such as BatchNorm and RNN. These variables are not updated by gradients, but by update ops.

We need to run the update ops to update these internal states of the layers.



--- END ISSUE ---

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

Path: `elasticdl/examples/mnist/mnist.py`
Content:
```
1 import tensorflow as tf
2 tf.enable_eager_execution()
3 
4 import os
5 import argparse
6 import numpy as np
7 from worker.worker import Worker
8 
9 
10 class MnistModel(tf.keras.Model):
11     def __init__(self, channel_last=True):
12         super(MnistModel, self).__init__(name='mnist_model')
13         if channel_last:
14             self._reshape = tf.keras.layers.Reshape((28, 28, 1))
15         else:
16             self._reshape = tf.keras.layers.Reshape((1, 28, 28))
17         self._conv1 = tf.keras.layers.Conv2D(
18             32, kernel_size=(3, 3), activation='relu')
19         self._conv2 = tf.keras.layers.Conv2D(
20             64, kernel_size=(3, 3), activation='relu')
21         self._batch_norm = tf.keras.layers.BatchNormalization()
22         self._maxpooling = tf.keras.layers.MaxPooling2D(
23             pool_size=(2, 2))
24         self._dropout = tf.keras.layers.Dropout(0.25)
25         self._flatten = tf.keras.layers.Flatten()
26         self._dense = tf.keras.layers.Dense(10)
27 
28     def call(self, inputs):
29         x = self._reshape(inputs)
30         x = self._conv1(x)
31         x = self._conv2(x)
32         x = self._batch_norm(x)
33         x = self._maxpooling(x)
34         x = self._dropout(x)
35         x = self._flatten(x)
36         x = self._dense(x)
37         return x
38 
39     @staticmethod
40     def input_shapes():
41         return (1, 28, 28)
42 
43     @staticmethod
44     def input_names():
45         return ['image']
46 
47     @staticmethod
48     def loss(output, labels):
49         return tf.reduce_mean(
50             tf.nn.sparse_softmax_cross_entropy_with_logits(
51                 logits=output, labels=labels['label']))
52 
53     @staticmethod
54     def optimizer(lr=0.1):
55         return tf.train.GradientDescentOptimizer(lr)
56 
57     @staticmethod
58     def input_fn(records):
59         image_list = []
60         label_list = []
61         # deserialize
62         for r in records:
63             parsed = np.frombuffer(r, dtype="uint8")
64             label = parsed[-1]
65             image = np.resize(parsed[:-1], new_shape=(28, 28))
66             image = image.astype(np.float32)
67             image /= 255
68             label = label.astype(np.int32)
69             image_list.append(image)
70             label_list.append(label)
71 
72         # batching
73         batch_size = len(image_list)
74         images = np.concatenate(image_list, axis=0)
75         images = np.reshape(images, (batch_size, 28, 28))
76         labels = np.array(label_list)
77         return {'image': images, 'label': labels}
78
```
--- 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/elasticdl/examples/mnist/mnist.py b/elasticdl/examples/mnist/mnist.py
--- a/elasticdl/examples/mnist/mnist.py
+++ b/elasticdl/examples/mnist/mnist.py
@@ -25,13 +25,14 @@
         self._flatten = tf.keras.layers.Flatten()
         self._dense = tf.keras.layers.Dense(10)
 
-    def call(self, inputs):
+    def call(self, inputs, training=False):
         x = self._reshape(inputs)
         x = self._conv1(x)
         x = self._conv2(x)
-        x = self._batch_norm(x)
+        x = self._batch_norm(x, training=training)
         x = self._maxpooling(x)
-        x = self._dropout(x)
+        if training:
+            x = self._dropout(x, training=training)
         x = self._flatten(x)
         x = self._dense(x)
         return x
diff --git a/elasticdl/worker/worker.py b/elasticdl/worker/worker.py
--- a/elasticdl/worker/worker.py
+++ b/elasticdl/worker/worker.py
@@ -114,7 +114,7 @@
                                     inputs.append(batch_input_data[input_name])
                                 if len(inputs) == 1:
                                     inputs = inputs[0]
-                                outputs = self._model.call(inputs)
+                                outputs = self._model.call(inputs, training=True)
                                 loss = self._model.loss(outputs, batch_input_data)
 
                                 # TODO:  Add regularization loss if any,
@@ -164,7 +164,7 @@
                                 inputs.append(data[input_name])
                             if len(inputs) == 1:
                                 inputs = inputs[0]
-                            outputs = self._model.call(inputs)
+                            outputs = self._model.call(inputs, training=True)
                             loss = self._model.loss(outputs, data)
 
                             # Add regularization loss if any.

{"golden_diff": "diff --git a/elasticdl/examples/mnist/mnist.py b/elasticdl/examples/mnist/mnist.py\n--- a/elasticdl/examples/mnist/mnist.py\n+++ b/elasticdl/examples/mnist/mnist.py\n@@ -25,13 +25,14 @@\n         self._flatten = tf.keras.layers.Flatten()\n         self._dense = tf.keras.layers.Dense(10)\n \n-    def call(self, inputs):\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)\n+        x = self._batch_norm(x, training=training)\n         x = self._maxpooling(x)\n-        x = self._dropout(x)\n+        if training:\n+            x = self._dropout(x, training=training)\n         x = self._flatten(x)\n         x = self._dense(x)\n         return x\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@@ -114,7 +114,7 @@\n                                     inputs.append(batch_input_data[input_name])\n                                 if len(inputs) == 1:\n                                     inputs = inputs[0]\n-                                outputs = self._model.call(inputs)\n+                                outputs = self._model.call(inputs, training=True)\n                                 loss = self._model.loss(outputs, batch_input_data)\n \n                                 # TODO:  Add regularization loss if any,\n@@ -164,7 +164,7 @@\n                                 inputs.append(data[input_name])\n                             if len(inputs) == 1:\n                                 inputs = inputs[0]\n-                            outputs = self._model.call(inputs)\n+                            outputs = self._model.call(inputs, training=True)\n                             loss = self._model.loss(outputs, data)\n \n                             # Add regularization loss if any.\n", "issue": "non-trainable variables update in training\nSome keras layers have non-trainable variables, such as BatchNorm and RNN. These variables are not updated by gradients, but by update ops.\r\n\r\nWe need to run the update ops to update these internal states of the layers.\r\n\r\n\n", "before_files": [{"content": "import tensorflow as tf\nassert tf.executing_eagerly()\n\nfrom google.protobuf import empty_pb2\nfrom tensorflow.python.ops import math_ops\nfrom proto import master_pb2_grpc\nfrom proto import master_pb2\nfrom common.ndarray import ndarray_to_tensor, tensor_to_ndarray\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_cls,\n                 channel=None,\n                 max_retrain_num=DEFAULT_MAX_MINIBATCH_RETRAIN_NUM):\n        \"\"\"\n        Arguments:\n            model_cls: A class 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        self._model = model_cls()\n        self._model.build(model_cls.input_shapes())\n\n        self._input_fn = model_cls.input_fn \n        self._opt_fn = model_cls.optimizer\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._model.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)\n                                loss = self._model.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._model.input_names():\n                                inputs.append(data[input_name])\n                            if len(inputs) == 1:\n                                inputs = inputs[0]\n                            outputs = self._model.call(inputs)\n                            loss = self._model.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"}, {"content": "import tensorflow as tf\ntf.enable_eager_execution()\n\nimport os\nimport argparse\nimport numpy as np\nfrom worker.worker import Worker\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):\n        x = self._reshape(inputs)\n        x = self._conv1(x)\n        x = self._conv2(x)\n        x = self._batch_norm(x)\n        x = self._maxpooling(x)\n        x = self._dropout(x)\n        x = self._flatten(x)\n        x = self._dense(x)\n        return x\n\n    @staticmethod\n    def input_shapes():\n        return (1, 28, 28)\n\n    @staticmethod\n    def input_names():\n        return ['image']\n\n    @staticmethod\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\n    @staticmethod\n    def optimizer(lr=0.1):\n        return tf.train.GradientDescentOptimizer(lr)\n\n    @staticmethod\n    def 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": "elasticdl/examples/mnist/mnist.py"}], "after_files": [{"content": "import tensorflow as tf\nassert tf.executing_eagerly()\n\nfrom google.protobuf import empty_pb2\nfrom tensorflow.python.ops import math_ops\nfrom proto import master_pb2_grpc\nfrom proto import master_pb2\nfrom common.ndarray import ndarray_to_tensor, tensor_to_ndarray\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_cls,\n                 channel=None,\n                 max_retrain_num=DEFAULT_MAX_MINIBATCH_RETRAIN_NUM):\n        \"\"\"\n        Arguments:\n            model_cls: A class 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        self._model = model_cls()\n        self._model.build(model_cls.input_shapes())\n\n        self._input_fn = model_cls.input_fn \n        self._opt_fn = model_cls.optimizer\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._model.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._model.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._model.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._model.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"}, {"content": "import tensorflow as tf\ntf.enable_eager_execution()\n\nimport os\nimport argparse\nimport numpy as np\nfrom worker.worker import Worker\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    @staticmethod\n    def input_shapes():\n        return (1, 28, 28)\n\n    @staticmethod\n    def input_names():\n        return ['image']\n\n    @staticmethod\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\n    @staticmethod\n    def optimizer(lr=0.1):\n        return tf.train.GradientDescentOptimizer(lr)\n\n    @staticmethod\n    def 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": "elasticdl/examples/mnist/mnist.py"}]}