Main class of super resolution model
| 10 | import tensorflow as tf |
| 11 | |
| 12 | class EDSRModel(tf.keras.Model): |
| 13 | """ |
| 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 | """ |
| 42 | prediction function |
| 43 | """ |
| 44 | # Adding dummy dimension using tf.expand_dims and converting to float32 using tf.cast |
| 45 | out = tf.cast(tf.expand_dims(inputs, axis=0), tf.float32) |
| 46 | # Passing low resolution image to model |
| 47 | super_resolution_img = self(out, training=False) |
| 48 | # Clips the tensor from min(0) to max(255) |
| 49 | super_resolution_img = tf.clip_by_value(super_resolution_img, 0, 255) |
| 50 | # Rounds the values of a tensor to the nearest integer |
| 51 | super_resolution_img = tf.round(super_resolution_img) |
| 52 | # Removes dimensions of size 1 from the shape of a tensor and converting to uint8 |
| 53 | super_resolution_img = tf.squeeze( |
| 54 | tf.cast(super_resolution_img, tf.uint8), axis=0 |
| 55 | ) |
| 56 | return super_resolution_img |
| 57 | |
| 58 | # Residual Block |
| 59 | def resblock(inputs): |