Saves the gradient of embedding activations ops in a graph collection.
(activations_op, grad_wrt_activations)
| 164 | |
| 165 | @ops.RegisterGradient("TPUEmbeddingActivations") |
| 166 | def _embedding_activations_grad(activations_op, grad_wrt_activations): |
| 167 | """Saves the gradient of embedding activations ops in a graph collection.""" |
| 168 | g = ops.get_default_graph() |
| 169 | table_id = activations_op.get_attr("table_id") |
| 170 | lookup_id = activations_op.get_attr("lookup_id") |
| 171 | table_gradients = g.get_collection_ref( |
| 172 | "tpu_embedding_gradients_table_%d" % table_id) |
| 173 | |
| 174 | if not table_gradients: |
| 175 | raise RuntimeError( |
| 176 | "Gradients for TPUEmbedding have been generated in non-training mode." |
| 177 | "This is not expected. Consider putting your Optimizer.minimize code " |
| 178 | "behind the training mode condition check. For Estimator, you can " |
| 179 | "do \n\n" |
| 180 | " if mode == tf.estimator.ModeKeys.TRAIN:\n" |
| 181 | " train_op = opt.minimize(loss)\n" |
| 182 | "\n") |
| 183 | |
| 184 | table_gradients[lookup_id] = array_ops.identity(grad_wrt_activations) |
| 185 | return [ |
| 186 | # RegisterGradient requires that value be returned for all inputs. Since |
| 187 | # the first argument (tpu_gradient_variable_{table_name}) has shape [1], |
| 188 | # we will return zeros(shape=[1]). The actual gradient w.r.t. the |
| 189 | # embedding activations (grad_wrt_activations) has the same shape as the |
| 190 | # activations returned by embedding_activations. |
| 191 | array_ops.zeros(arg.shape, dtype=dtypes.float32) |
| 192 | for arg in activations_op.inputs |
| 193 | ] |
| 194 | |
| 195 | |
| 196 | def infeed_dequeue(dtype, shape, name=None): |
nothing calls this directly
no test coverage detected