Module degann.networks.topology.tf_densenet

Classes

class TensorflowDenseNet (input_size: int = 2, block_size: list = None, output_size: int = 10, activation_func: str = 'linear', weight=<keras.src.initializers.random_initializers.RandomUniform object>, biases=<keras.src.initializers.random_initializers.RandomUniform object>, is_debug: bool = False, **kwargs)

A model grouping layers into an object with training/inference features.

There are three ways to instantiate a Model:

With the "Functional API"

You start from Input, you chain layer calls to specify the model's forward pass, and finally you create your model from inputs and outputs:

inputs = keras.Input(shape=(37,))
x = keras.layers.Dense(32, activation="relu")(inputs)
outputs = keras.layers.Dense(5, activation="softmax")(x)
model = keras.Model(inputs=inputs, outputs=outputs)

Note: Only dicts, lists, and tuples of input tensors are supported. Nested inputs are not supported (e.g. lists of list or dicts of dict).

A new Functional API model can also be created by using the intermediate tensors. This enables you to quickly extract sub-components of the model.

Example:

inputs = keras.Input(shape=(None, None, 3))
processed = keras.layers.RandomCrop(width=128, height=128)(inputs)
conv = keras.layers.Conv2D(filters=32, kernel_size=3)(processed)
pooling = keras.layers.GlobalAveragePooling2D()(conv)
feature = keras.layers.Dense(10)(pooling)

full_model = keras.Model(inputs, feature)
backbone = keras.Model(processed, conv)
activations = keras.Model(conv, feature)

Note that the backbone and activations models are not created with keras.Input objects, but with the tensors that originate from keras.Input objects. Under the hood, the layers and weights will be shared across these models, so that user can train the full_model, and use backbone or activations to do feature extraction. The inputs and outputs of the model can be nested structures of tensors as well, and the created models are standard Functional API models that support all the existing APIs.

By subclassing the Model class

In that case, you should define your layers in __init__() and you should implement the model's forward pass in call().

class MyModel(keras.Model):
    def __init__(self):
        super().__init__()
        self.dense1 = keras.layers.Dense(32, activation="relu")
        self.dense2 = keras.layers.Dense(5, activation="softmax")

    def call(self, inputs):
        x = self.dense1(inputs)
        return self.dense2(x)

model = MyModel()

If you subclass Model, you can optionally have a training argument (boolean) in call(), which you can use to specify a different behavior in training and inference:

class MyModel(keras.Model):
    def __init__(self):
        super().__init__()
        self.dense1 = keras.layers.Dense(32, activation="relu")
        self.dense2 = keras.layers.Dense(5, activation="softmax")
        self.dropout = keras.layers.Dropout(0.5)

    def call(self, inputs, training=False):
        x = self.dense1(inputs)
        x = self.dropout(x, training=training)
        return self.dense2(x)

model = MyModel()

Once the model is created, you can config the model with losses and metrics with model.compile(), train the model with model.fit(), or use the model to do prediction with model.predict().

With the Sequential class

In addition, keras.Sequential is a special case of model where the model is purely a stack of single-input, single-output layers.

model = keras.Sequential([
    keras.Input(shape=(None, None, 3)),
    keras.layers.Conv2D(filters=32, kernel_size=3),
])

Expand source code

class TensorflowDenseNet(tf.keras.Model):
    def __init__(
        self,
        input_size: int = 2,
        block_size: list = None,
        output_size: int = 10,
        activation_func: str = "linear",
        weight=keras.initializers.RandomUniform(minval=-1, maxval=1),
        biases=keras.initializers.RandomUniform(minval=-1, maxval=1),
        is_debug: bool = False,
        **kwargs,
    ):
        decorator_params: List[Optional[Dict]] = [None]
        if "decorator_params" in kwargs.keys():
            decorator_params = kwargs.get("decorator_params")
            kwargs.pop("decorator_params")
        else:
            decorator_params = [None]

        if (
            isinstance(decorator_params, list)
            and len(decorator_params) == 1
            and decorator_params[0] is None
            or decorator_params is None
        ):
            decorator_params = [None] * (len(block_size) + 1)

        if (
            isinstance(decorator_params, list)
            and len(decorator_params) == 1
            and decorator_params[0] is not None
        ):
            decorator_params = decorator_params * (len(block_size) + 1)

        super(TensorflowDenseNet, self).__init__(**kwargs)
        self.blocks: List[TensorflowDense] = []

        if not isinstance(activation_func, list):
            activation_func = [activation_func] * (len(block_size) + 1)
        if len(block_size) != 0:
            self.blocks.append(
                layer_creator.create_dense(
                    input_size,
                    block_size[0],
                    activation=activation_func[0],
                    weight=weight,
                    bias=biases,
                    is_debug=is_debug,
                    name=f"MyDense0",
                    decorator_params=decorator_params[0],
                )
            )
            for i in range(1, len(block_size)):
                self.blocks.append(
                    layer_creator.create_dense(
                        block_size[i - 1],
                        block_size[i],
                        activation=activation_func[i],
                        weight=weight,
                        bias=biases,
                        is_debug=is_debug,
                        name=f"MyDense{i}",
                        decorator_params=decorator_params[i],
                    )
                )
            last_block_size = block_size[-1]
        else:
            last_block_size = input_size

        self.out_layer = layer_creator.create_dense(
            last_block_size,
            output_size,
            activation=activation_func[-1],
            weight=weight,
            bias=biases,
            is_debug=is_debug,
            name=f"OutLayerMyDense",
            decorator_params=decorator_params[-1],
        )

        self.activation_funcs = activation_func
        self.weight_initializer = weight
        self.bias_initializer = biases
        self.input_size = input_size
        self.block_size = block_size
        self.output_size = output_size
        self.trained_time = {"train_time": 0.0, "epoch_time": [], "predict_time": 0}

    def custom_compile(
        self,
        rate=1e-2,
        optimizer="SGD",
        loss_func="MeanSquaredError",
        metric_funcs=None,
        run_eagerly=False,
    ):
        """
        Configures the model for training

        Parameters
        ----------
        rate: float
            learning rate for optimizer
        optimizer: str
            name of optimizer
        loss_func: str
            name of loss function
        metric_funcs: list[str]
            list with metric function names
        run_eagerly: bool

        Returns
        -------

        """
        opt = optimizers.get_optimizer(optimizer)(learning_rate=rate)
        loss = losses.get_loss(loss_func)
        m = [metrics.get_metric(metric) for metric in metric_funcs]
        self.compile(
            optimizer=opt,
            loss=loss,
            metrics=m,
            run_eagerly=run_eagerly,
        )

    def call(self, inputs, **kwargs):
        """
        Obtaining a neural network response on the input data vector
        Parameters
        ----------
        inputs
        kwargs

        Returns
        -------

        """
        x = inputs
        for layer in self.blocks:
            x = layer(x, **kwargs)
        return self.out_layer(x, **kwargs)

    def train_step(self, data):
        """
        Custom train step from tensorflow tutorial

        Parameters
        ----------
        data: tuple
            Pair of x and y (or dataset)
        Returns
        -------

        """
        # Unpack the data. Its structure depends on your model and
        # on what you pass to `fit()`.
        x, y = data
        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)  # Forward pass
            # Compute the loss value
            # (the loss function is configured in `compile()`)
            loss = self.compute_loss(y=y, y_pred=y_pred)

        # Compute gradients
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)
        # Update weights
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))
        # Update metrics (includes the metric that tracks the loss)
        for metric in self.metrics:
            if metric.name == "loss":
                metric.update_state(loss)
            else:
                metric.update_state(y, y_pred)
        # Return a dict mapping metric names to current value
        return {m.name: m.result() for m in self.metrics}

    def set_name(self, new_name):
        self._name = new_name

    def __str__(self):
        res = f"IModel {self.name}\n"
        for layer in self.blocks:
            res += str(layer)
        res += str(self.out_layer)
        return res

    def to_dict(self, **kwargs):
        """
        Export neural network to dictionary

        Parameters
        ----------
        kwargs

        Returns
        -------

        """
        res = {
            "net_type": "MyDense",
            "name": self._name,
            "input_size": self.input_size,
            "block_size": self.block_size,
            "output_size": self.output_size,
            "layer": [],
            "out_layer": self.out_layer.to_dict(),
        }

        for i, layer in enumerate(self.blocks):
            res["layer"].append(layer.to_dict())

        return res

    @classmethod
    def from_layers(
        cls,
        input_size: int,
        block_size: List[int],
        output_size: int,
        layers: List[TensorflowDense],
        **kwargs,
    ):
        """
        Restore neural network from list of layers
        Parameters
        ----------
        input_size
        block_size
        output_size
        layers
        kwargs

        Returns
        -------

        """
        res = cls(
            input_size=input_size,
            block_size=block_size,
            output_size=output_size,
            **kwargs,
        )

        for layer_num in range(len(res.blocks)):
            res.blocks[layer_num] = layers[layer_num]

        return res

    def from_dict(self, config, **kwargs):
        """
        Restore neural network from dictionary of params
        Parameters
        ----------
        config
        kwargs

        Returns
        -------

        """
        input_size = config["input_size"]
        block_size = config["block_size"]
        output_size = config["output_size"]

        self.block_size = list(block_size)
        self.input_size = input_size
        self.output_size = output_size

        layers: List[TensorflowDense] = []
        for layer_config in config["layer"]:
            layers.append(layer_creator.from_dict(layer_config))

        self.blocks: List[TensorflowDense] = []
        for layer_num in range(len(layers)):
            self.blocks.append(layers[layer_num])

        self.out_layer = layer_creator.from_dict(config["out_layer"])

    def export_to_cpp(
        self, path: str, array_type: str = "[]", path_to_compiler: str = None, **kwargs
    ) -> None:
        """
        Export neural network as feedforward function on c++

        Parameters
        ----------
        path: str
            path to file with name, without extension
        array_type: str
            c-style or cpp-style ("[]" or "vector")
        path_to_compiler: str
            path to c/c++ compiler
        kwargs

        Returns
        -------

        """
        res = """
#include <cmath>
#include <vector>

#define max(x, y) ((x < y) ? y : x)

        \n"""

        config = self.to_dict(**kwargs)

        input_size = self.input_size
        output_size = self.output_size
        blocks = self.block_size
        reverse = False
        layers = config["layer"] + [config["out_layer"]]

        comment = f"// This function takes {input_size} elements array and returns {output_size} elements array\n"
        signature = f""
        start_func = "{\n"
        end_func = "}\n"
        transform_input_vector = ""
        transform_output_array = ""
        return_stat = "return answer;\n"

        creator_1d: Callable[
            [str, int, Optional[list]], str
        ] = cpp_utils.array1d_creator("float")
        creator_heap_1d: Callable[[str, int], str] = cpp_utils.array1d_heap_creator(
            "float"
        )
        creator_2d: Callable[
            [str, int, int, Optional[list]], str
        ] = cpp_utils.array2d_creator("float")
        if array_type == "[]":
            signature = f"float* feedforward(float x_array[])\n"

        if array_type == "vector":
            signature = f"std::vector<float> feedforward(std::vector<float> x)\n"

            transform_input_vector = cpp_utils.transform_1dvector_to_array(
                "float", input_size, "x", "x_array"
            )
            transform_output_array = cpp_utils.transform_1darray_to_vector(
                "float", output_size, "answer_vector", "answer"
            )
            return_stat = "return answer_vector;\n"

        create_layers = ""
        create_layers += creator_1d(f"layer_0", input_size, initial_value=0)
        for i, size in enumerate(blocks):
            create_layers += creator_1d(f"layer_{i + 1}", size, initial_value=0)
        create_layers += creator_1d(
            f"layer_{len(blocks) + 1}", output_size, initial_value=0
        )
        create_layers += cpp_utils.copy_1darray_to_array(
            input_size, "x_array", "layer_0"
        )

        create_weights = ""
        for i, layer_dict in enumerate(layers):
            create_weights += creator_2d(
                f"weight_{i}_{i + 1}",
                layer_dict[LAYER_DICT_NAMES["inp_size"]],
                layer_dict[LAYER_DICT_NAMES["shape"]],
                layer_dict[LAYER_DICT_NAMES["weights"]],
                reverse=reverse,
            )

        fill_weights = ""

        create_biases = ""
        for i, layer_dict in enumerate(layers):
            create_biases += creator_1d(
                f"bias_{i + 1}",
                layer_dict[LAYER_DICT_NAMES["shape"]],
                layer_dict[LAYER_DICT_NAMES["bias"]],
            )

        fill_biases = ""
        feed_forward_cycles = ""
        for i, layer_dict in enumerate(layers):
            left_size = layer_dict[
                LAYER_DICT_NAMES["inp_size"]
            ]  # if i != 0 else input_size
            right_size = layer_dict[LAYER_DICT_NAMES["shape"]]
            act_func = layer_dict[LAYER_DICT_NAMES["activation"]]
            decorator_params = layer_dict.get(LAYER_DICT_NAMES["decorator_params"])
            feed_forward_cycles += cpp_utils.feed_forward_step(
                f"layer_{i}",
                left_size,
                f"layer_{i + 1}",
                right_size,
                f"weight_{i}_{i + 1}",
                f"bias_{i + 1}",
                act_func,
            )

        move_result = creator_heap_1d("answer", output_size)
        move_result += cpp_utils.copy_1darray_to_array(
            output_size, f"layer_{len(blocks) + 1}", "answer"
        )

        res += comment
        res += signature
        res += start_func
        res += transform_input_vector
        res += create_layers
        res += create_weights
        res += fill_weights
        res += create_biases
        res += fill_biases
        res += feed_forward_cycles
        res += move_result
        res += transform_output_array
        res += return_stat
        res += end_func

        header_res = f"""
        #ifndef {path[0].upper() + path[1:]}_hpp
        #define {path[0].upper() + path[1:]}_hpp
        #include "{path}.cpp"

        {comment}
        {signature};

        #endif /* {path[0].upper() + path[1:]}_hpp */

                """

        with open(path + ".cpp", "w") as f:
            f.write(res)

        with open(path + ".hpp", "w") as f:
            f.write(header_res)

        if path_to_compiler is not None:
            os.system(path_to_compiler + " -c -Ofast " + path + ".cpp")

    @property
    def get_activations(self) -> List:
        """
        Get list of activations functions for each layer

        Returns
        -------
        activation: list
        """
        return [layer.get_activation for layer in self.blocks]

Ancestors

• keras.src.models.model.Model
• keras.src.backend.tensorflow.trainer.TensorFlowTrainer
• keras.src.trainers.trainer.Trainer
• keras.src.layers.layer.Layer
• keras.src.backend.tensorflow.layer.TFLayer
• keras.src.backend.tensorflow.trackable.KerasAutoTrackable
• tensorflow.python.trackable.autotrackable.AutoTrackable
• tensorflow.python.trackable.base.Trackable
• keras.src.ops.operation.Operation
• keras.src.saving.keras_saveable.KerasSaveable

Static methods

def from_layers(input_size: int, block_size: List[int], output_size: int, layers: List[TensorflowDense], **kwargs)

Restore neural network from list of layers Parameters

---

input_size

 

block_size

 

output_size

 

layers

 

kwargs

 

Returns

Expand source code

@classmethod
def from_layers(
    cls,
    input_size: int,
    block_size: List[int],
    output_size: int,
    layers: List[TensorflowDense],
    **kwargs,
):
    """
    Restore neural network from list of layers
    Parameters
    ----------
    input_size
    block_size
    output_size
    layers
    kwargs

    Returns
    -------

    """
    res = cls(
        input_size=input_size,
        block_size=block_size,
        output_size=output_size,
        **kwargs,
    )

    for layer_num in range(len(res.blocks)):
        res.blocks[layer_num] = layers[layer_num]

    return res

Instance variables

var get_activations : List

Get list of activations functions for each layer

Returns

activation : list

 

Expand source code

@property
def get_activations(self) -> List:
    """
    Get list of activations functions for each layer

    Returns
    -------
    activation: list
    """
    return [layer.get_activation for layer in self.blocks]

Methods

def call(self, inputs, **kwargs)

Obtaining a neural network response on the input data vector Parameters

---

inputs

 

kwargs

 

Returns

Expand source code

def call(self, inputs, **kwargs):
    """
    Obtaining a neural network response on the input data vector
    Parameters
    ----------
    inputs
    kwargs

    Returns
    -------

    """
    x = inputs
    for layer in self.blocks:
        x = layer(x, **kwargs)
    return self.out_layer(x, **kwargs)

def custom_compile(self, rate=0.01, optimizer='SGD', loss_func='MeanSquaredError', metric_funcs=None, run_eagerly=False)

Configures the model for training

Parameters

rate : float

learning rate for optimizer

optimizer : str

name of optimizer

loss_func : str

name of loss function

metric_funcs : list[str]

list with metric function names

run_eagerly : bool

 

Returns

Expand source code

def custom_compile(
    self,
    rate=1e-2,
    optimizer="SGD",
    loss_func="MeanSquaredError",
    metric_funcs=None,
    run_eagerly=False,
):
    """
    Configures the model for training

    Parameters
    ----------
    rate: float
        learning rate for optimizer
    optimizer: str
        name of optimizer
    loss_func: str
        name of loss function
    metric_funcs: list[str]
        list with metric function names
    run_eagerly: bool

    Returns
    -------

    """
    opt = optimizers.get_optimizer(optimizer)(learning_rate=rate)
    loss = losses.get_loss(loss_func)
    m = [metrics.get_metric(metric) for metric in metric_funcs]
    self.compile(
        optimizer=opt,
        loss=loss,
        metrics=m,
        run_eagerly=run_eagerly,
    )

def export_to_cpp(self, path: str, array_type: str = '[]', path_to_compiler: str = None, **kwargs) ‑> None

Export neural network as feedforward function on c++

Parameters

path : str

path to file with name, without extension

array_type : str

c-style or cpp-style ("[]" or "vector")

path_to_compiler : str

path to c/c++ compiler

kwargs

 

Returns

Expand source code

    def export_to_cpp(
        self, path: str, array_type: str = "[]", path_to_compiler: str = None, **kwargs
    ) -> None:
        """
        Export neural network as feedforward function on c++

        Parameters
        ----------
        path: str
            path to file with name, without extension
        array_type: str
            c-style or cpp-style ("[]" or "vector")
        path_to_compiler: str
            path to c/c++ compiler
        kwargs

        Returns
        -------

        """
        res = """
#include <cmath>
#include <vector>

#define max(x, y) ((x < y) ? y : x)

        \n"""

        config = self.to_dict(**kwargs)

        input_size = self.input_size
        output_size = self.output_size
        blocks = self.block_size
        reverse = False
        layers = config["layer"] + [config["out_layer"]]

        comment = f"// This function takes {input_size} elements array and returns {output_size} elements array\n"
        signature = f""
        start_func = "{\n"
        end_func = "}\n"
        transform_input_vector = ""
        transform_output_array = ""
        return_stat = "return answer;\n"

        creator_1d: Callable[
            [str, int, Optional[list]], str
        ] = cpp_utils.array1d_creator("float")
        creator_heap_1d: Callable[[str, int], str] = cpp_utils.array1d_heap_creator(
            "float"
        )
        creator_2d: Callable[
            [str, int, int, Optional[list]], str
        ] = cpp_utils.array2d_creator("float")
        if array_type == "[]":
            signature = f"float* feedforward(float x_array[])\n"

        if array_type == "vector":
            signature = f"std::vector<float> feedforward(std::vector<float> x)\n"

            transform_input_vector = cpp_utils.transform_1dvector_to_array(
                "float", input_size, "x", "x_array"
            )
            transform_output_array = cpp_utils.transform_1darray_to_vector(
                "float", output_size, "answer_vector", "answer"
            )
            return_stat = "return answer_vector;\n"

        create_layers = ""
        create_layers += creator_1d(f"layer_0", input_size, initial_value=0)
        for i, size in enumerate(blocks):
            create_layers += creator_1d(f"layer_{i + 1}", size, initial_value=0)
        create_layers += creator_1d(
            f"layer_{len(blocks) + 1}", output_size, initial_value=0
        )
        create_layers += cpp_utils.copy_1darray_to_array(
            input_size, "x_array", "layer_0"
        )

        create_weights = ""
        for i, layer_dict in enumerate(layers):
            create_weights += creator_2d(
                f"weight_{i}_{i + 1}",
                layer_dict[LAYER_DICT_NAMES["inp_size"]],
                layer_dict[LAYER_DICT_NAMES["shape"]],
                layer_dict[LAYER_DICT_NAMES["weights"]],
                reverse=reverse,
            )

        fill_weights = ""

        create_biases = ""
        for i, layer_dict in enumerate(layers):
            create_biases += creator_1d(
                f"bias_{i + 1}",
                layer_dict[LAYER_DICT_NAMES["shape"]],
                layer_dict[LAYER_DICT_NAMES["bias"]],
            )

        fill_biases = ""
        feed_forward_cycles = ""
        for i, layer_dict in enumerate(layers):
            left_size = layer_dict[
                LAYER_DICT_NAMES["inp_size"]
            ]  # if i != 0 else input_size
            right_size = layer_dict[LAYER_DICT_NAMES["shape"]]
            act_func = layer_dict[LAYER_DICT_NAMES["activation"]]
            decorator_params = layer_dict.get(LAYER_DICT_NAMES["decorator_params"])
            feed_forward_cycles += cpp_utils.feed_forward_step(
                f"layer_{i}",
                left_size,
                f"layer_{i + 1}",
                right_size,
                f"weight_{i}_{i + 1}",
                f"bias_{i + 1}",
                act_func,
            )

        move_result = creator_heap_1d("answer", output_size)
        move_result += cpp_utils.copy_1darray_to_array(
            output_size, f"layer_{len(blocks) + 1}", "answer"
        )

        res += comment
        res += signature
        res += start_func
        res += transform_input_vector
        res += create_layers
        res += create_weights
        res += fill_weights
        res += create_biases
        res += fill_biases
        res += feed_forward_cycles
        res += move_result
        res += transform_output_array
        res += return_stat
        res += end_func

        header_res = f"""
        #ifndef {path[0].upper() + path[1:]}_hpp
        #define {path[0].upper() + path[1:]}_hpp
        #include "{path}.cpp"

        {comment}
        {signature};

        #endif /* {path[0].upper() + path[1:]}_hpp */

                """

        with open(path + ".cpp", "w") as f:
            f.write(res)

        with open(path + ".hpp", "w") as f:
            f.write(header_res)

        if path_to_compiler is not None:
            os.system(path_to_compiler + " -c -Ofast " + path + ".cpp")

def from_dict(self, config, **kwargs)

Restore neural network from dictionary of params Parameters

---

config

 

kwargs

 

Returns

Expand source code

def from_dict(self, config, **kwargs):
    """
    Restore neural network from dictionary of params
    Parameters
    ----------
    config
    kwargs

    Returns
    -------

    """
    input_size = config["input_size"]
    block_size = config["block_size"]
    output_size = config["output_size"]

    self.block_size = list(block_size)
    self.input_size = input_size
    self.output_size = output_size

    layers: List[TensorflowDense] = []
    for layer_config in config["layer"]:
        layers.append(layer_creator.from_dict(layer_config))

    self.blocks: List[TensorflowDense] = []
    for layer_num in range(len(layers)):
        self.blocks.append(layers[layer_num])

    self.out_layer = layer_creator.from_dict(config["out_layer"])

def set_name(self, new_name)

Expand source code

def set_name(self, new_name):
    self._name = new_name

def to_dict(self, **kwargs)

Export neural network to dictionary

Parameters

kwargs

 

Returns

Expand source code

def to_dict(self, **kwargs):
    """
    Export neural network to dictionary

    Parameters
    ----------
    kwargs

    Returns
    -------

    """
    res = {
        "net_type": "MyDense",
        "name": self._name,
        "input_size": self.input_size,
        "block_size": self.block_size,
        "output_size": self.output_size,
        "layer": [],
        "out_layer": self.out_layer.to_dict(),
    }

    for i, layer in enumerate(self.blocks):
        res["layer"].append(layer.to_dict())

    return res

def train_step(self, data)

Custom train step from tensorflow tutorial

Parameters

data : tuple

Pair of x and y (or dataset)

Returns

Expand source code

def train_step(self, data):
    """
    Custom train step from tensorflow tutorial

    Parameters
    ----------
    data: tuple
        Pair of x and y (or dataset)
    Returns
    -------

    """
    # Unpack the data. Its structure depends on your model and
    # on what you pass to `fit()`.
    x, y = data
    with tf.GradientTape() as tape:
        y_pred = self(x, training=True)  # Forward pass
        # Compute the loss value
        # (the loss function is configured in `compile()`)
        loss = self.compute_loss(y=y, y_pred=y_pred)

    # Compute gradients
    trainable_vars = self.trainable_variables
    gradients = tape.gradient(loss, trainable_vars)
    # Update weights
    self.optimizer.apply_gradients(zip(gradients, trainable_vars))
    # Update metrics (includes the metric that tracks the loss)
    for metric in self.metrics:
        if metric.name == "loss":
            metric.update_state(loss)
        else:
            metric.update_state(y, y_pred)
    # Return a dict mapping metric names to current value
    return {m.name: m.result() for m in self.metrics}