enable visualizations for tensorboard tensorboard is a visualization tool provided with tensorflow a tensorflow installation is required to use this callback this callback logs events for tensorboard including metrics summary plots training graph visualization weight histograms sampled profiling when used in model evaluate or regular validation in addition to epoch summaries there will be a summary that records evaluation metrics vs model optimizer iterations written the metric names will be prepended with evaluation with model optimizer iterations being the step in the visualized tensorboard if you have installed tensorflow with pip you should be able to launch tensorboard from the command line tensorboard logdirpathtoyourlogs you can find more information about tensorboard herehttps www tensorflow orggetstartedsummariesandtensorboard args logdir the path of the directory where to save the log files to be parsed by tensorboard e g logdir os path joinworkingdir logs this directory should not be reused by any other callbacks histogramfreq frequency in epochs at which to compute weight histograms for the layers of the model if set to 0 histograms won t be computed validation data or split must be specified for histogram visualizations writegraph not supported at this time whether to visualize the graph in tensorboard note that the log file can become quite large when writegraph is set to true writeimages whether to write model weights to visualize as image in tensorboard writestepspersecond whether to log the training steps per second into tensorboard this supports both epoch and batch frequency logging updatefreq batch or epoch or integer when using epoch writes the losses and metrics to tensorboard after every epoch if using an integer let s say 1000 all metrics and losses including custom ones added by model compile will be logged to tensorboard every 1000 batches batch is a synonym for 1 meaning that they will be written every batch note however that writing too frequently to tensorboard can slow down your training especially when used with distribution strategies as it will incur additional synchronization overhead batchlevel summary writing is also available via trainstep override please see tensorboard scalars tutorial https www tensorflow orgtensorboardscalarsandkerasbatchlevellogging noqa e501 for more details profilebatch not supported at this time profile the batches to sample compute characteristics profilebatch must be a nonnegative integer or a tuple of integers a pair of positive integers signify a range of batches to profile by default profiling is disabled embeddingsfreq frequency in epochs at which embedding layers will be visualized if set to 0 embeddings won t be visualized embeddingsmetadata dictionary which maps embedding layer names to the filename of a file in which to save metadata for the embedding layer in case the same metadata file is to be used for all embedding layers a single filename can be passed examples basic usage python tensorboardcallback keras callbacks tensorboardlogdir logs model fitxtrain ytrain epochs2 callbackstensorboardcallback then run the tensorboard command to view the visualizations custom batchlevel summaries in a subclassed model python class mymodelkeras model def buildself self dense keras layers dense10 def callself x outputs self densex tf summary histogram outputs outputs return outputs model mymodel model compile sgd mse make sure to set updatefreqn to log a batchlevel summary every n batches in addition to any tf summary contained in model call metrics added in model compile will be logged every n batches tbcallback keras callbacks tensorboard logs updatefreq1 model fitxtrain ytrain callbackstbcallback custom batchlevel summaries in a functional api model python def mysummaryx tf summary histogram x x return x inputs keras input10 x keras layers dense10inputs outputs keras layers lambdamysummaryx model keras modelinputs outputs model compile sgd mse make sure to set updatefreqn to log a batchlevel summary every n batches in addition to any tf summary contained in model call metrics added in model compile will be logged every n batches tbcallback keras callbacks tensorboard logs updatefreq1 model fitxtrain ytrain callbackstbcallback profiling python profile a single batch e g the 5th batch tensorboardcallback keras callbacks tensorboard logdir logs profilebatch5 model fitxtrain ytrain epochs2 callbackstensorboardcallback profile a range of batches e g from 10 to 20 tensorboardcallback keras callbacks tensorboard logdir logs profilebatch10 20 model fitxtrain ytrain epochs2 callbackstensorboardcallback todo profiling not available in jaxtorch lazily initialized in order to avoid creating event files when not needed used to restore any existing summarywriter after training ends sets keras model and writes graph if specified self model model self logwritedir self logdir self traindir os path joinself logwritedir train self trainstep 0 self valdir os path joinself logwritedir validation self valstep 0 self writers resets writers self shouldwritetraingraph false if self writegraph self writekerasmodelsummary self shouldwritetraingraph true if self embeddingsfreq self configureembeddings property def summaryself if self summarymodule is none import tensorflow summary as summary self summarymodule summary return self summarymodule property def trainwriterself if train not in self writers self writerstrain self summary createfilewriter self traindir return self writerstrain property def valwriterself if val not in self writers self writersval self summary createfilewriter self valdir return self writersval def writekerasmodeltraingraphself if the trainfunction is a tf function we can write out a graph todob243822285 use variablecreationfn directly writes keras graph network summary to tensorboard with self trainwriter asdefault with self summary recordiftrue if self model class name functional or self model class name sequential kerasmodelsummarykeras self model step0 def configureembeddingsself embeddings are always the first layer so this naming should be consistent in any keras models checkpoints sets the default writer for custom batchlevel summaries if self updatefreq epoch return def shouldrecord return step self updatefreq 0 summarycontext writer asdefaultstep self summary recordifshouldrecord self prevsummarystate appendsummarycontext summarycontext0 enter summarycontext1 enter def popwriterself see pushwriter for the content of the previouscontext which is pair of context validate profilebatch value and set the range of batches to profile sets values of startbatch and stopbatch attributes specifying the start and stop batch to profile setting profilebatch0 disables profiling args profilebatch the range of batches to profile should be a nonnegative integer or a comma separated string of pair of positive integers a pair of positive integers signify a range of batches to profile raises valueerror if profilebatch is not an integer or a comma separated pair of positive integers support legacy way of specifying start stop or start as str true when the profiler was successfully started by this callback we track the status here to make sure callbacks do not interfere with each other the callback will only stop the profiler it started warm up and improve the profiling accuracy true when a trace is running setting profilebatch0 disables profiling only call batch hooks when tracing or writestepspersecond are enabled logs isn t necessarily always a dict keeps track of epoch for profiling runs metrics and histogram summaries at epoch end self logepochmetricsepoch logs if self histogramfreq and epoch self histogramfreq 0 self logweightsepoch if self embeddingsfreq and epoch self embeddingsfreq 0 self logembeddingsepoch def starttraceself self summary traceongraphtrue profilerfalse self startprofilerlogdirself logdir self istracing true def stoptraceself batchnone todob126388999 remove step info in the summary name writes epoch metrics out as scalar summaries args epoch int the global step to use for tensorboard logs dict keys are scalar summary names values are scalars logs the weights of the model to tensorboard with self trainwriter asdefault with self summary recordiftrue for layer in self model layers for weight in layer weights weightname weight name replace add a suffix to prevent summary tag name collision histogramweightname weightname histogram self summary histogram histogramweightname weight stepepoch if self writeimages add a suffix to prevent summary tag name collision imageweightname weightname image self logweightasimage weight imageweightname epoch self trainwriter flush def logweightasimageself weight weightname epoch switch to channelsfirst to display every kernel as a separate image not possible to handle 3d convnets etc starts the profiler if currently inactive args logdir directory where profiler results will be saved profiler errors should not be fatal stops the profiler if currently active args save whether to save the profiler results to tensorboard profiler errors should not be fatal writes a keras model as json to as a summary writing the keras model configuration allows the tensorboard graph plugin to render a conceptual graph as opposed to graph of ops in case the model fails to serialize as json it ignores and returns false args name a name for this summary the summary tag used for tensorboard will be this name prefixed by any active name scopes data a keras model to write step explicit int64castable monotonic step value for this summary if omitted this defaults to tf summary experimental getstep which must not be none returns true on success or false if no summary was written because no default summary writer was available raises valueerror if a default writer exists but no step was provided and tf summary experimental getstep is none hard coding a plugin name please refer to gotbpluginnamehardcode for the rationale version number 1 an exception should not break a model code enable visualizations for tensorboard tensorboard is a visualization tool provided with tensorflow a tensorflow installation is required to use this callback this callback logs events for tensorboard including metrics summary plots training graph visualization weight histograms sampled profiling when used in model evaluate or regular validation in addition to epoch summaries there will be a summary that records evaluation metrics vs model optimizer iterations written the metric names will be prepended with evaluation with model optimizer iterations being the step in the visualized tensorboard if you have installed tensorflow with pip you should be able to launch tensorboard from the command line tensorboard logdir path_to_your_logs you can find more information about tensorboard here https www tensorflow org get_started summaries_and_tensorboard args log_dir the path of the directory where to save the log files to be parsed by tensorboard e g log_dir os path join working_dir logs this directory should not be reused by any other callbacks histogram_freq frequency in epochs at which to compute weight histograms for the layers of the model if set to 0 histograms won t be computed validation data or split must be specified for histogram visualizations write_graph not supported at this time whether to visualize the graph in tensorboard note that the log file can become quite large when write_graph is set to true write_images whether to write model weights to visualize as image in tensorboard write_steps_per_second whether to log the training steps per second into tensorboard this supports both epoch and batch frequency logging update_freq batch or epoch or integer when using epoch writes the losses and metrics to tensorboard after every epoch if using an integer let s say 1000 all metrics and losses including custom ones added by model compile will be logged to tensorboard every 1000 batches batch is a synonym for 1 meaning that they will be written every batch note however that writing too frequently to tensorboard can slow down your training especially when used with distribution strategies as it will incur additional synchronization overhead batch level summary writing is also available via train_step override please see tensorboard scalars tutorial https www tensorflow org tensorboard scalars_and_keras batch level_logging noqa e501 for more details profile_batch not supported at this time profile the batch es to sample compute characteristics profile_batch must be a non negative integer or a tuple of integers a pair of positive integers signify a range of batches to profile by default profiling is disabled embeddings_freq frequency in epochs at which embedding layers will be visualized if set to 0 embeddings won t be visualized embeddings_metadata dictionary which maps embedding layer names to the filename of a file in which to save metadata for the embedding layer in case the same metadata file is to be used for all embedding layers a single filename can be passed examples basic usage python tensorboard_callback keras callbacks tensorboard log_dir logs model fit x_train y_train epochs 2 callbacks tensorboard_callback then run the tensorboard command to view the visualizations custom batch level summaries in a subclassed model python class mymodel keras model def build self _ self dense keras layers dense 10 def call self x outputs self dense x tf summary histogram outputs outputs return outputs model mymodel model compile sgd mse make sure to set update_freq n to log a batch level summary every n batches in addition to any tf summary contained in model call metrics added in model compile will be logged every n batches tb_callback keras callbacks tensorboard logs update_freq 1 model fit x_train y_train callbacks tb_callback custom batch level summaries in a functional api model python def my_summary x tf summary histogram x x return x inputs keras input 10 x keras layers dense 10 inputs outputs keras layers lambda my_summary x model keras model inputs outputs model compile sgd mse make sure to set update_freq n to log a batch level summary every n batches in addition to any tf summary contained in model call metrics added in model compile will be logged every n batches tb_callback keras callbacks tensorboard logs update_freq 1 model fit x_train y_train callbacks tb_callback profiling python profile a single batch e g the 5th batch tensorboard_callback keras callbacks tensorboard log_dir logs profile_batch 5 model fit x_train y_train epochs 2 callbacks tensorboard_callback profile a range of batches e g from 10 to 20 tensorboard_callback keras callbacks tensorboard log_dir logs profile_batch 10 20 model fit x_train y_train epochs 2 callbacks tensorboard_callback todo profiling not available in jax torch lazily initialized in order to avoid creating event files when not needed used to restore any existing summarywriter after training ends sets keras model and writes graph if specified resets writers writes keras model train_function graph to tensorboard if the train_function is a tf function we can write out a graph todo b 243822285 use _variable_creation_fn directly writes keras graph network summary to tensorboard configure the projector for embeddings embeddings are always the first layer so this naming should be consistent in any keras models checkpoints sets the default writer for custom batch level summaries pops the current writer see _push_writer for the content of the previous_context which is pair of context validate profile_batch value and set the range of batches to profile sets values of _start_batch and _stop_batch attributes specifying the start and stop batch to profile setting profile_batch 0 disables profiling args profile_batch the range of batches to profile should be a non negative integer or a comma separated string of pair of positive integers a pair of positive integers signify a range of batches to profile raises valueerror if profile_batch is not an integer or a comma separated pair of positive integers support legacy way of specifying start stop or start as str true when the profiler was successfully started by this callback we track the status here to make sure callbacks do not interfere with each other the callback will only stop the profiler it started warm up and improve the profiling accuracy true when a trace is running setting profile_batch 0 disables profiling only call batch hooks when tracing or write_steps_per_second are enabled logs isn t necessarily always a dict keeps track of epoch for profiling runs metrics and histogram summaries at epoch end logs the trace graph to tensorboard todo b 126388999 remove step info in the summary name writes epoch metrics out as scalar summaries args epoch int the global step to use for tensorboard logs dict keys are scalar summary names values are scalars remove val_ prefix logs the weights of the model to tensorboard add a suffix to prevent summary tag name collision add a suffix to prevent summary tag name collision logs a weight as a tensorboard image bias case dense layer kernel case convnet case switch to channels_first to display every kernel as a separate image not possible to handle 3d convnets etc starts the profiler if currently inactive args logdir directory where profiler results will be saved profiler errors should not be fatal stops the profiler if currently active args save whether to save the profiler results to tensorboard profiler errors should not be fatal writes a keras model as json to as a summary writing the keras model configuration allows the tensorboard graph plugin to render a conceptual graph as opposed to graph of ops in case the model fails to serialize as json it ignores and returns false args name a name for this summary the summary tag used for tensorboard will be this name prefixed by any active name scopes data a keras model to write step explicit int64 castable monotonic step value for this summary if omitted this defaults to tf summary experimental get_step which must not be none returns true on success or false if no summary was written because no default summary writer was available raises valueerror if a default writer exists but no step was provided and tf summary experimental get_step is none hard coding a plugin name please refer to go tb plugin name hardcode for the rationale version number 1 an exception should not break a model code
import logging
import os
import sys
import time
import warnings

import tree

from keras import backend
from keras import ops
from keras.api_export import keras_export
from keras.callbacks.callback import Callback
from keras.layers import Embedding
from keras.optimizers import Optimizer
from keras.utils import file_utils


@keras_export("keras.callbacks.TensorBoard")
class TensorBoard(Callback):
    

    def __init__(
        self,
        log_dir="logs",
        histogram_freq=0,
        write_graph=True,
        write_images=False,
        write_steps_per_second=False,
        update_freq="epoch",
        profile_batch=0,
        embeddings_freq=0,
        embeddings_metadata=None,
    ):
        super().__init__()

        self.log_dir = str(log_dir)
        self.histogram_freq = histogram_freq
        self.write_graph = write_graph
        self.write_images = write_images
        self.write_steps_per_second = write_steps_per_second
        self.update_freq = 1 if update_freq == "batch" else update_freq
        self.embeddings_freq = embeddings_freq
        self.embeddings_metadata = embeddings_metadata
        if profile_batch and backend.backend() != "tensorflow":
            
            raise ValueError(
                "Profiling is not yet available with the "
                f"{backend.backend()} backend. Please open a PR "
                "if you'd like to add this feature. Received: "
                f"profile_batch={profile_batch} (must be 0)"
            )
        self._init_profile_batch(profile_batch)
        self._global_train_batch = 0
        self._previous_epoch_iterations = 0
        self._train_accumulated_time = 0
        self._batch_start_time = 0
        self._summary_module = None

        
        
        self._writers = {}

        
        self._prev_summary_state = []

    def set_model(self, model):
        
        self._model = model
        self._log_write_dir = self.log_dir

        self._train_dir = os.path.join(self._log_write_dir, "train")
        self._train_step = 0

        self._val_dir = os.path.join(self._log_write_dir, "validation")
        self._val_step = 0

        self._writers = {}  

        self._should_write_train_graph = False
        if self.write_graph:
            self._write_keras_model_summary()
            self._should_write_train_graph = True
        if self.embeddings_freq:
            self._configure_embeddings()

    @property
    def summary(self):
        if self._summary_module is None:
            import tensorflow.summary as summary

            self._summary_module = summary
        return self._summary_module

    @property
    def _train_writer(self):
        if "train" not in self._writers:
            self._writers["train"] = self.summary.create_file_writer(
                self._train_dir
            )
        return self._writers["train"]

    @property
    def _val_writer(self):
        if "val" not in self._writers:
            self._writers["val"] = self.summary.create_file_writer(
                self._val_dir
            )
        return self._writers["val"]

    def _write_keras_model_train_graph(self):
        
        with self._train_writer.as_default():
            with self.summary.record_if(True):
                train_fn = self.model.train_function
                
                
                if hasattr(train_fn, "function_spec"):
                    
                    if hasattr(train_fn, "_concrete_stateful_fn"):
                        self.summary.graph(train_fn._concrete_stateful_fn.graph)
                    else:
                        self.summary.graph(
                            train_fn._concrete_variable_creation_fn.graph
                        )

    def _write_keras_model_summary(self):
        
        with self._train_writer.as_default():
            with self.summary.record_if(True):
                if (
                    self.model.__class__.__name__ == "Functional"
                    or self.model.__class__.__name__ == "Sequential"
                ):
                    keras_model_summary("keras", self.model, step=0)

    def _configure_embeddings(self):
        
        from google.protobuf import text_format
        from tensorboard.plugins import projector

        config = projector.ProjectorConfig()
        for layer in self.model.layers:
            if isinstance(layer, Embedding):
                embedding = config.embeddings.add()
                
                
                name = (
                    "layer_with_weights-0/embeddings/.ATTRIBUTES/VARIABLE_VALUE"
                )
                embedding.tensor_name = name

                if self.embeddings_metadata is not None:
                    if isinstance(self.embeddings_metadata, str):
                        embedding.metadata_path = self.embeddings_metadata
                    else:
                        if layer.name in self.embeddings_metadata.keys():
                            embedding.metadata_path = (
                                self.embeddings_metadata.pop(layer.name)
                            )

        if self.embeddings_metadata and not isinstance(
            self.embeddings_metadata, str
        ):
            raise ValueError(
                "Unrecognized `Embedding` layer names passed to "
                "`keras.callbacks.TensorBoard` `embeddings_metadata` "
                f"argument: {self.embeddings_metadata.keys()}"
            )

        config_pbtxt = text_format.MessageToString(config)
        path = os.path.join(self._log_write_dir, "projector_config.pbtxt")
        with file_utils.File(path, "w") as f:
            f.write(config_pbtxt)

    def _push_writer(self, writer, step):
        
        if self.update_freq == "epoch":
            return

        def should_record():
            return step % self.update_freq == 0

        summary_context = (
            writer.as_default(step),
            self.summary.record_if(should_record),
        )
        self._prev_summary_state.append(summary_context)
        summary_context[0].__enter__()
        summary_context[1].__enter__()

    def _pop_writer(self):
        
        if self.update_freq == "epoch":
            return

        
        
        previous_context = self._prev_summary_state.pop()
        previous_context[1].__exit__(*sys.exc_info())
        previous_context[0].__exit__(*sys.exc_info())

    def _close_writers(self):
        for writer in self._writers.values():
            writer.close()

    def _init_profile_batch(self, profile_batch):
        
        profile_batch_error_message = (
            "profile_batch must be a non-negative integer or "
            "2-tuple of positive "
            "integers. A pair of positive integers "
            "signifies a range of batches "
            f"to profile. Found: {profile_batch}"
        )

        
        if isinstance(profile_batch, str):
            profile_batch = str(profile_batch).split(",")
            profile_batch = tree.map_structure(int, profile_batch)

        if isinstance(profile_batch, int):
            self._start_batch = profile_batch
            self._stop_batch = profile_batch
        elif (
            isinstance(profile_batch, (tuple, list)) and len(profile_batch) == 2
        ):
            self._start_batch, self._stop_batch = profile_batch
        else:
            raise ValueError(profile_batch_error_message)

        if self._start_batch < 0 or self._stop_batch < self._start_batch:
            raise ValueError(profile_batch_error_message)

        
        
        
        self._profiler_started = False
        if self._start_batch > 0:
            
            self._start_profiler(logdir="")
            self._stop_profiler(save=False)
        
        self._is_tracing = False

        
        self._should_trace = not (
            self._start_batch == 0 and self._stop_batch == 0
        )

    def on_train_begin(self, logs=None):
        self._global_train_batch = 0
        self._previous_epoch_iterations = 0
        self._push_writer(self._train_writer, self._train_step)

    def on_train_end(self, logs=None):
        self._pop_writer()

        if self._is_tracing:
            self._stop_trace()

        self._close_writers()

    def on_test_begin(self, logs=None):
        self._push_writer(self._val_writer, self._val_step)

    def on_test_end(self, logs=None):
        if self.model.optimizer and hasattr(self.model.optimizer, "iterations"):
            with self.summary.record_if(True), self._val_writer.as_default():
                for name, value in logs.items():
                    self.summary.scalar(
                        "evaluation_" + name + "_vs_iterations",
                        value,
                        step=self.model.optimizer.iterations,
                    )
        self._pop_writer()

    def _implements_train_batch_hooks(self):
        
        
        return self._should_trace or self.write_steps_per_second

    def on_train_batch_begin(self, batch, logs=None):
        self._global_train_batch += 1
        if self.write_steps_per_second:
            self._batch_start_time = time.time()
        if not self._should_trace:
            return

        if self._global_train_batch == self._start_batch:
            self._start_trace()

    def on_train_batch_end(self, batch, logs=None):
        if self._should_write_train_graph:
            self._write_keras_model_train_graph()
            self._should_write_train_graph = False
        if self.write_steps_per_second:
            batch_run_time = time.time() - self._batch_start_time
            self.summary.scalar(
                "batch_steps_per_second",
                1.0 / batch_run_time,
                step=self._train_step,
            )

        
        if isinstance(logs, dict):
            for name, value in logs.items():
                self.summary.scalar(
                    "batch_" + name, value, step=self._train_step
                )

        if not self._should_trace:
            return

        if self._is_tracing and self._global_train_batch >= self._stop_batch:
            self._stop_trace()

    def on_epoch_begin(self, epoch, logs=None):
        
        if self.write_steps_per_second:
            self._previous_epoch_iterations = self.model.optimizer.iterations
            self._epoch_start_time = time.time()

    def on_epoch_end(self, epoch, logs=None):
        
        self._log_epoch_metrics(epoch, logs)

        if self.histogram_freq and epoch % self.histogram_freq == 0:
            self._log_weights(epoch)

        if self.embeddings_freq and epoch % self.embeddings_freq == 0:
            self._log_embeddings(epoch)

    def _start_trace(self):
        self.summary.trace_on(graph=True, profiler=False)
        self._start_profiler(logdir=self.log_dir)
        self._is_tracing = True

    def _stop_trace(self, batch=None):
        
        if batch is None:
            batch = self._stop_batch
        with self._train_writer.as_default():
            with self.summary.record_if(True):
                
                self.summary.trace_export(name="batch_%d" % batch, step=batch)
        self._stop_profiler()
        self._is_tracing = False

    def _collect_learning_rate(self, logs):
        if isinstance(self.model.optimizer, Optimizer):
            logs["learning_rate"] = float(
                ops.convert_to_numpy(self.model.optimizer.learning_rate)
            )
        return logs

    def _compute_steps_per_second(self):
        current_iteration = self.model.optimizer.iterations
        time_since_epoch_begin = time.time() - self._epoch_start_time
        current_iteration = ops.convert_to_tensor(current_iteration, "float32")
        self._previous_epoch_iterations = ops.convert_to_tensor(
            self._previous_epoch_iterations, "float32"
        )
        time_since_epoch_begin = ops.convert_to_tensor(
            time_since_epoch_begin, "float32"
        )

        steps_per_second = (
            current_iteration - self._previous_epoch_iterations
        ) / time_since_epoch_begin
        return float(steps_per_second)

    def _log_epoch_metrics(self, epoch, logs):
        
        if not logs:
            return

        train_logs = {k: v for k, v in logs.items() if not k.startswith("val_")}
        val_logs = {k: v for k, v in logs.items() if k.startswith("val_")}
        train_logs = self._collect_learning_rate(train_logs)
        if self.write_steps_per_second:
            train_logs["steps_per_second"] = self._compute_steps_per_second()

        with self.summary.record_if(True):
            if train_logs:
                with self._train_writer.as_default():
                    for name, value in train_logs.items():
                        self.summary.scalar("epoch_" + name, value, step=epoch)
            if val_logs:
                with self._val_writer.as_default():
                    for name, value in val_logs.items():
                        name = name[4:]  
                        self.summary.scalar("epoch_" + name, value, step=epoch)

    def _log_weights(self, epoch):
        
        with self._train_writer.as_default():
            with self.summary.record_if(True):
                for layer in self.model.layers:
                    for weight in layer.weights:
                        weight_name = weight.name.replace(":", "_")
                        
                        histogram_weight_name = weight_name + "/histogram"
                        self.summary.histogram(
                            histogram_weight_name, weight, step=epoch
                        )
                        if self.write_images:
                            
                            
                            image_weight_name = weight_name + "/image"
                            self._log_weight_as_image(
                                weight, image_weight_name, epoch
                            )
                self._train_writer.flush()

    def _log_weight_as_image(self, weight, weight_name, epoch):
        
        w_img = ops.squeeze(weight)
        shape = w_img.shape
        if len(shape) == 1:  
            w_img = ops.reshape(w_img, [1, shape[0], 1, 1])
        elif len(shape) == 2:  
            if shape[0] > shape[1]:
                w_img = ops.transpose(w_img)
                shape = w_img.shape
            w_img = ops.reshape(w_img, [1, shape[0], shape[1], 1])
        elif len(shape) == 3:  
            if backend.image_data_format() == "channels_last":
                
                
                w_img = ops.transpose(w_img, [2, 0, 1])
                shape = w_img.shape
            w_img = ops.reshape(w_img, [shape[0], shape[1], shape[2], 1])

        w_img = backend.convert_to_numpy(w_img)
        shape = w_img.shape
        
        if len(shape) == 4 and shape[-1] in [1, 3, 4]:
            self.summary.image(weight_name, w_img, step=epoch)

    def _log_embeddings(self, epoch):
        embeddings_ckpt = os.path.join(
            self._log_write_dir,
            "train",
            f"keras_embedding.ckpt-{epoch}.weights.h5",
        )
        self.model.save_weights(embeddings_ckpt)

    def _start_profiler(self, logdir):
        
        if self._profiler_started:
            return
        try:
            backend.tensorboard.start_trace(logdir)
            self._profiler_started = True
        except Exception as e:
            
            logging.error("Failed to start profiler: %s", e)

    def _stop_profiler(self, save=True):
        
        if not self._profiler_started:
            return
        try:
            backend.tensorboard.stop_trace(save=save)
        except Exception as e:
            
            logging.error("Failed to stop profiler: %s", e)
        finally:
            self._profiler_started = False


def keras_model_summary(name, data, step=None):
    
    import tensorflow.summary as summary
    from tensorflow.compat.v1 import SummaryMetadata

    summary_metadata = SummaryMetadata()
    
    
    summary_metadata.plugin_data.plugin_name = "graph_keras_model"
    
    summary_metadata.plugin_data.content = b"1"

    try:
        json_string = data.to_json()
    except Exception as exc:
        
        warnings.warn(f"Model failed to serialize as JSON. Ignoring... {exc}")
        return False

    with summary.experimental.summary_scope(
        name, "graph_keras_model", [data, step]
    ) as (tag, _):
        return summary.write(
            tag=tag, tensor=json_string, step=step, metadata=summary_metadata
        )