MCPcopy Create free account
hub / github.com/CandleLabAI/PCBSegClassNet / train_step

Method train_step

src/models/edsr.py:16–38  ·  view source on GitHub ↗

forward pass function

(self, data)

Source from the content-addressed store, hash-verified

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 """

Callers

nothing calls this directly

Calls

no outgoing calls

Tested by

no test coverage detected