| 53 | |
| 54 | |
| 55 | class _ModelWithOptimizer(util.Checkpoint): |
| 56 | |
| 57 | def __init__(self): |
| 58 | self.dense = core.Dense(1) |
| 59 | self.optimizer = adam.Adam(0.01) |
| 60 | |
| 61 | @def_function.function( |
| 62 | input_signature=(tensor_spec.TensorSpec([None, 2], dtypes.float32), |
| 63 | tensor_spec.TensorSpec([None], dtypes.float32))) |
| 64 | def call(self, x, y): |
| 65 | with backprop.GradientTape() as tape: |
| 66 | loss = math_ops.reduce_mean((self.dense(x) - y) ** 2.) |
| 67 | trainable_variables = self.dense.trainable_variables |
| 68 | gradients = tape.gradient(loss, trainable_variables) |
| 69 | self.optimizer.apply_gradients(zip(gradients, trainable_variables)) |
| 70 | return {"loss": loss} |
| 71 | |
| 72 | |
| 73 | def _import_and_infer( |
no outgoing calls