forward pass function
(self, data)
| 14 | Main class of super resolution model |
| 15 | """ |
| 16 | def train_step(self, data): |
| 17 | """ |
| 18 | forward pass function |
| 19 | """ |
| 20 | # Unpack the data. Its structure depends on your model and |
| 21 | # on what you pass to `fit()`. |
| 22 | inp, target = data |
| 23 | |
| 24 | with tf.GradientTape() as tape: |
| 25 | y_pred = self(inp, training=True) # Forward pass |
| 26 | # Compute the loss value |
| 27 | # (the loss function is configured in `compile()`) |
| 28 | loss = self.compiled_loss(target, y_pred, regularization_losses=self.losses) |
| 29 | |
| 30 | # Compute gradients |
| 31 | trainable_vars = self.trainable_variables |
| 32 | gradients = tape.gradient(loss, trainable_vars) |
| 33 | # Update weights |
| 34 | self.optimizer.apply_gradients(zip(gradients, trainable_vars)) |
| 35 | # Update metrics (includes the metric that tracks the loss) |
| 36 | self.compiled_metrics.update_state(target, y_pred) |
| 37 | # Return a dict mapping metric names to current value |
| 38 | return {m.name: m.result() for m in self.metrics} |
| 39 | |
| 40 | def predict_step(self, inputs): |
| 41 | """ |
nothing calls this directly
no outgoing calls
no test coverage detected