require torch2 1 0 to enable dynamo since it includes many improvementsfixes to torch compile todo eventually we want to get rid of this when torch is upgraded to 2 1 from 2 0 1 in g3 compute predictions call torch nn module zerograd to clear the leftover gradients for the weights from the previous train step compute gradients call torch tensor backward on the loss to compute gradients for the weights update weights runs a single training step on a batch of data data data0 return self trainstepdata if self shouldtorchcompile self trainfunction torch compileonestepondata else self trainfunction onestepondata def maketestfunctionself forcefalse if self testfunction is not none and not force return self testfunction if self stepsperexecution 1 raise valueerror stepsperexecution must be 1 with the pytorch backend freceived stepsperexecutionself stepsperexecution def onestepondatadata runs a predict test step on a batch of data data data0 with torch nograd return self predictstepdata if self shouldtorchcompile self predictfunction torch compileonestepondata else self predictfunction onestepondata def symbolicbuildself databatch modelunbuilt not alllayer built for layer in self flattenlayers compilemetricsunbuilt self compilemetrics is not none and not self compilemetrics built if modelunbuilt or compilemetricsunbuilt create symbolic tensors matching an input batch def tosymbolicinputv if istensorv return kerastensorv shape standardizedtypev dtype return v databatch tree mapstructuretosymbolicinput databatch x y sampleweight dataadapterutils unpackxysampleweightdatabatch build all model state with backend computeoutputspec try ypred backend computeoutputspecself x except exception as e raise runtimeerror unable to automatically build the model please build it yourself before calling fitevaluatepredict a model is built when its variables have been created and its self built attribute is true usually calling the model on a batch of data is the right way to build it n exception encountered n f e if compilemetricsunbuilt build all metric state with backend computeoutputspec backend computeoutputspec self computemetrics x y ypred sampleweightsampleweight if self optimizer is not none and not self optimizer built build optimizer self optimizer buildself trainablevariables self postbuild tracebackutils filtertraceback def fit self xnone ynone batchsizenone epochs1 verboseauto callbacksnone validationsplit0 0 validationdatanone shuffletrue classweightnone sampleweightnone initialepoch0 stepsperepochnone validationstepsnone validationbatchsizenone validationfreq1 if not self compiled raise valueerror you must call compile before calling fit todo respect compiled trainable state self evalepochiterator none if validationsplit and validationdata is none create the validation data using the training data only supported for tfnumpyjax arrays todo support torch tensors for validation data x y sampleweight validationdata dataadapterutils trainvalidationsplit x y sampleweight validationsplitvalidationsplit if validationdata valx valy valsampleweight dataadapterutils unpackxysampleweightvalidationdata create an iterator that yields batches for one epoch epochiterator epochiterator xx yy sampleweightsampleweight batchsizebatchsize stepsperepochstepsperepoch shuffleshuffle classweightclassweight stepsperexecutionself stepsperexecution needsbuilding not alllayer built for layer in self flattenlayers or not self optimizer built or self compilemetrics is not none and not self compilemetrics built if needsbuilding build the model on one batch of data for data in epochiterator enumerateepoch databatch data0 self symbolicbuilddatabatch break container that configures and calls callbacks if not isinstancecallbacks callbacksmodule callbacklist callbacks callbacksmodule callbacklist callbacks addhistorytrue addprogbarverbose 0 verboseverbose epochsepochs stepsepochiterator numbatches modelself self stoptraining false self maketrainfunction callbacks ontrainbegin for epoch in rangeinitialepoch epochs self resetmetrics callbacks onepochbeginepoch switch the torch module to training mode inform torch layers to do training behavior in case the user did not use self training when implementing a custom layer with torch layers self train for step data in epochiterator enumerateepoch callbacks callbacks ontrainbatchbeginstep logs self trainfunctiondata callbacks callbacks ontrainbatchendstep self pythonifylogslogs if self stoptraining break override with model metrics instead of last step logs epochlogs self getmetricsresult switch the torch module back to testing mode self eval run validation if validationdata and self shouldevalepoch validationfreq create epochiterator for evaluation and cache it if getattrself evalepochiterator none is none self evalepochiterator epochiterator xvalx yvaly sampleweightvalsampleweight batchsizevalidationbatchsize or batchsize stepsperexecutionself stepsperexecution stepsperepochvalidationsteps shufflefalse vallogs self evaluate xvalx yvaly sampleweightvalsampleweight batchsizevalidationbatchsize or batchsize stepsvalidationsteps callbackscallbacks returndicttrue usecachedevaldatasettrue vallogs val name val for name val in vallogs items epochlogs updatevallogs callbacks onepochendepoch epochlogs traininglogs epochlogs if self stoptraining break if isinstanceself optimizer optimizersmodule optimizer and epochs 0 self optimizer finalizevariablevaluesself trainableweights if evalepochiterator exists delete it after all epochs are done if getattrself evalepochiterator none is not none del self evalepochiterator callbacks ontrainendlogstraininglogs return self history tracebackutils filtertraceback def evaluate self xnone ynone batchsizenone verboseauto sampleweightnone stepsnone callbacksnone returndictfalse kwargs todo respect compiled trainable state usecachedevaldataset kwargs popusecachedevaldataset false if kwargs raise valueerrorfarguments not recognized kwargs if usecachedevaldataset epochiterator self evalepochiterator else create an iterator that yields batches of inputtarget data epochiterator epochiterator xx yy sampleweightsampleweight batchsizebatchsize stepsperepochsteps shufflefalse stepsperexecutionself stepsperexecution if not alllayer built for layer in self flattenlayers build the model on one batch of data for data in epochiterator enumerateepoch databatch data0 self symbolicbuilddatabatch break container that configures and calls callbacks if not isinstancecallbacks callbacksmodule callbacklist callbacks callbacksmodule callbacklist callbacks addhistorytrue addprogbarverbose 0 verboseverbose epochs1 stepsepochiterator numbatches modelself switch the torch module back to testing mode self eval self maketestfunction self stopevaluating false callbacks ontestbegin logs none self resetmetrics for step data in epochiterator enumerateepoch callbacks ontestbatchbeginstep logs self testfunctiondata callbacks ontestbatchendstep self pythonifylogslogs if self stopevaluating break logs self getmetricsresult callbacks ontestendlogs if returndict return logs return self flattenmetricsinorderlogs tracebackutils filtertraceback def predict self x batchsizenone verboseauto stepsnone callbacksnone create an iterator that yields batches of input data epochiterator torchepochiterator xx batchsizebatchsize stepsperepochsteps shufflefalse stepsperexecutionself stepsperexecution container that configures and calls callbacks if not isinstancecallbacks callbacksmodule callbacklist callbacks callbacksmodule callbacklist callbacks addhistorytrue addprogbarverbose 0 verboseverbose epochs1 stepsepochiterator numbatches modelself def appendtooutputsbatchoutputs outputs if outputs is none outputs tree mapstructure lambda batchoutput batchoutput batchoutputs else tree mapstructureupto batchoutputs lambda output batchoutput output appendbatchoutput outputs batchoutputs return outputs switch the torch module back to testing mode self eval self makepredictfunction self stoppredicting false callbacks onpredictbegin outputs none for step data in epochiterator enumerateepoch callbacks onpredictbatchbeginstep batchoutputs self predictfunctiondata outputs appendtooutputsbatchoutputs outputs callbacks onpredictbatchendstep outputs batchoutputs if self stoppredicting break callbacks onpredictend outputs tree mapstructurebackend converttonumpy outputs return tree mapstructureuptobatchoutputs np concatenate outputs def trainonbatch self x ynone sampleweightnone classweightnone returndictfalse self assertcompilecalledtrainonbatch if classweight is not none if sampleweight is not none raise valueerror arguments sampleweight and classweight cannot be specified at the same time freceived sampleweightsampleweight fclassweightclassweight sampleweight dataadapterutils classweighttosampleweights y classweight data x y sampleweight maybe build model self symbolicbuilddata self maketrainfunction logs self trainfunctiondata logs tree mapstructurelambda x np arrayx logs if returndict return logs return self flattenmetricsinorderlogs def testonbatch self x ynone sampleweightnone returndictfalse self assertcompilecalledtestonbatch data x y sampleweight maybe build model self symbolicbuilddata self maketestfunction logs self testfunctiondata logs tree mapstructurelambda x np arrayx logs if returndict return logs return self flattenmetricsinorderlogs def predictonbatchself x self makepredictfunction batchoutputs self predictfunctionx batchoutputs tree mapstructure backend converttonumpy batchoutputs return batchoutputs class torchepochiteratorepochiterator epochiterator def getiteratorself returntypeauto if returntype auto and isinstance self dataadapter dataadapters torchdataloaderadapter return self dataadapter gettorchdataloader return super getiteratorreturntype require torch 2 1 0 to enable dynamo since it includes many improvements fixes to torch compile todo eventually we want to get rid of this when torch is upgraded to 2 1 from 2 0 1 in g3 compute predictions call torch nn module zero_grad to clear the leftover gradients for the weights from the previous train step compute gradients call torch tensor backward on the loss to compute gradients for the weights update weights runs a single training step on a batch of data runs a single test step on a batch of data runs a predict test step on a batch of data create symbolic tensors matching an input batch build all model state with backend compute_output_spec build all metric state with backend compute_output_spec build optimizer todo respect compiled trainable state create the validation data using the training data only supported for tf numpy jax arrays todo support torch tensors for validation data create an iterator that yields batches for one epoch build the model on one batch of data container that configures and calls callbacks switch the torch module to training mode inform torch layers to do training behavior in case the user did not use self training when implementing a custom layer with torch layers callbacks callbacks override with model metrics instead of last step logs switch the torch module back to testing mode run validation create epochiterator for evaluation and cache it if _eval_epoch_iterator exists delete it after all epochs are done todo respect compiled trainable state create an iterator that yields batches of input target data build the model on one batch of data container that configures and calls callbacks switch the torch module back to testing mode create an iterator that yields batches of input data container that configures and calls callbacks switch the torch module back to testing mode maybe build model maybe build model
import warnings

import numpy as np
import torch
import tree
from packaging.version import parse

from keras import backend
from keras import callbacks as callbacks_module
from keras import optimizers as optimizers_module
from keras.backend.common import standardize_dtype
from keras.backend.common.keras_tensor import KerasTensor
from keras.backend.torch.core import is_tensor
from keras.trainers import data_adapters
from keras.trainers import epoch_iterator
from keras.trainers import trainer as base_trainer
from keras.trainers.data_adapters import data_adapter_utils
from keras.trainers.epoch_iterator import EpochIterator
from keras.utils import traceback_utils


class TorchTrainer(base_trainer.Trainer):
    def __init__(self):
        super().__init__()
        self.train_function = None
        self.test_function = None
        self.predict_function = None

    def _should_torch_compile(self):
        
        
        
        
        if self.jit_compile and parse(torch.__version__) < parse("2.1.0"):
            warnings.warn(
                "Please upgrade to torch>=2.1.0 for `jit_compile=True` "
                "to take effect. Using `jit_compile=False`"
            )
            self.jit_compile = False

        return self.jit_compile

    def train_step(self, data):
        x, y, sample_weight = data_adapter_utils.unpack_x_y_sample_weight(data)

        
        if self._call_has_training_arg:
            y_pred = self(x, training=True)
        else:
            y_pred = self(x)

        
        
        self.zero_grad()

        loss = self.compute_loss(
            x=x, y=y, y_pred=y_pred, sample_weight=sample_weight
        )
        self._loss_tracker.update_state(loss)
        if self.optimizer is not None:
            loss = self.optimizer.scale_loss(loss)

        
        if self.trainable_weights:
            
            
            loss.backward()

            trainable_weights = self.trainable_weights[:]
            gradients = [v.value.grad for v in trainable_weights]

            
            with torch.no_grad():
                self.optimizer.apply(gradients, trainable_weights)
        else:
            warnings.warn("The model does not have any trainable weights.")

        return self.compute_metrics(x, y, y_pred, sample_weight=sample_weight)

    def test_step(self, data):
        (
            x,
            y,
            sample_weight,
        ) = data_adapter_utils.unpack_x_y_sample_weight(data)
        if self._call_has_training_arg:
            y_pred = self(x, training=False)
        else:
            y_pred = self(x)
        loss = self.compute_loss(
            x=x, y=y, y_pred=y_pred, sample_weight=sample_weight
        )
        self._loss_tracker.update_state(loss)
        return self.compute_metrics(x, y, y_pred, sample_weight=sample_weight)

    def predict_step(self, data):
        x, _, _ = data_adapter_utils.unpack_x_y_sample_weight(data)
        if self._call_has_training_arg:
            y_pred = self(x, training=False)
        else:
            y_pred = self(x)
        return y_pred

    def make_train_function(self, force=False):
        if self.train_function is not None and not force:
            return self.train_function

        if self.steps_per_execution > 1:
            raise ValueError(
                "`steps_per_execution` must be 1 with the PyTorch backend. "
                f"Received: steps_per_execution={self.steps_per_execution}"
            )

        def one_step_on_data(data):
            
            data = data[0]
            return self.train_step(data)

        if self._should_torch_compile():
            self.train_function = torch.compile(one_step_on_data)
        else:
            self.train_function = one_step_on_data

    def make_test_function(self, force=False):
        if self.test_function is not None and not force:
            return self.test_function

        if self.steps_per_execution > 1:
            raise ValueError(
                "`steps_per_execution` must be 1 with the PyTorch backend. "
                f"Received: steps_per_execution={self.steps_per_execution}"
            )

        def one_step_on_data(data):
            
            data = data[0]
            with torch.no_grad():
                return self.test_step(data)

        if self._should_torch_compile():
            self.test_function = torch.compile(one_step_on_data)
        else:
            self.test_function = one_step_on_data

    def make_predict_function(self, force=False):
        if self.predict_function is not None and not force:
            return self.predict_function

        if self.steps_per_execution > 1:
            raise ValueError(
                "`steps_per_execution` must be 1 with the PyTorch backend. "
                f"Received: steps_per_execution={self.steps_per_execution}"
            )

        def one_step_on_data(data):
            
            data = data[0]
            with torch.no_grad():
                return self.predict_step(data)

        if self._should_torch_compile():
            self.predict_function = torch.compile(one_step_on_data)
        else:
            self.predict_function = one_step_on_data

    def _symbolic_build(self, data_batch):
        model_unbuilt = not all(layer.built for layer in self._flatten_layers())
        compile_metrics_unbuilt = (
            self._compile_metrics is not None
            and not self._compile_metrics.built
        )
        if model_unbuilt or compile_metrics_unbuilt:
            

            def to_symbolic_input(v):
                if is_tensor(v):
                    return KerasTensor(v.shape, standardize_dtype(v.dtype))
                return v

            data_batch = tree.map_structure(to_symbolic_input, data_batch)
            (
                x,
                y,
                sample_weight,
            ) = data_adapter_utils.unpack_x_y_sample_weight(data_batch)
            
            try:
                y_pred = backend.compute_output_spec(self, x)
            except Exception as e:
                raise RuntimeError(
                    "Unable to automatically build the model. "
                    "Please build it yourself before calling "
                    "fit/evaluate/predict. "
                    "A model is 'built' when its variables have "
                    "been created and its `self.built` attribute "
                    "is True. Usually, calling the model on a batch "
                    "of data is the right way to build it.\n"
                    "Exception encountered:\n"
                    f"'{e}'"
                )
            if compile_metrics_unbuilt:
                
                backend.compute_output_spec(
                    self.compute_metrics,
                    x,
                    y,
                    y_pred,
                    sample_weight=sample_weight,
                )
        if self.optimizer is not None and not self.optimizer.built:
            
            self.optimizer.build(self.trainable_variables)
        self._post_build()

    @traceback_utils.filter_traceback
    def fit(
        self,
        x=None,
        y=None,
        batch_size=None,
        epochs=1,
        verbose="auto",
        callbacks=None,
        validation_split=0.0,
        validation_data=None,
        shuffle=True,
        class_weight=None,
        sample_weight=None,
        initial_epoch=0,
        steps_per_epoch=None,
        validation_steps=None,
        validation_batch_size=None,
        validation_freq=1,
    ):
        if not self.compiled:
            raise ValueError(
                "You must call `compile()` before calling `fit()`."
            )

        
        self._eval_epoch_iterator = None
        if validation_split and validation_data is None:
            
            
            
            (
                x,
                y,
                sample_weight,
            ), validation_data = data_adapter_utils.train_validation_split(
                (x, y, sample_weight), validation_split=validation_split
            )

        if validation_data:
            (
                val_x,
                val_y,
                val_sample_weight,
            ) = data_adapter_utils.unpack_x_y_sample_weight(validation_data)

        
        epoch_iterator = EpochIterator(
            x=x,
            y=y,
            sample_weight=sample_weight,
            batch_size=batch_size,
            steps_per_epoch=steps_per_epoch,
            shuffle=shuffle,
            class_weight=class_weight,
            steps_per_execution=self.steps_per_execution,
        )

        needs_building = (
            not all(layer.built for layer in self._flatten_layers())
            or not self.optimizer.built
            or (
                self._compile_metrics is not None
                and not self._compile_metrics.built
            )
        )
        if needs_building:
            
            for _, data in epoch_iterator.enumerate_epoch():
                data_batch = data[0]
                self._symbolic_build(data_batch)
                break

        
        if not isinstance(callbacks, callbacks_module.CallbackList):
            callbacks = callbacks_module.CallbackList(
                callbacks,
                add_history=True,
                add_progbar=verbose != 0,
                verbose=verbose,
                epochs=epochs,
                steps=epoch_iterator.num_batches,
                model=self,
            )

        self.stop_training = False
        self.make_train_function()
        callbacks.on_train_begin()

        for epoch in range(initial_epoch, epochs):
            self.reset_metrics()
            callbacks.on_epoch_begin(epoch)

            
            
            
            self.train()
            for step, data in epoch_iterator.enumerate_epoch():
                
                callbacks.on_train_batch_begin(step)

                logs = self.train_function(data)

                
                callbacks.on_train_batch_end(step, self._pythonify_logs(logs))
                if self.stop_training:
                    break

            
            epoch_logs = self.get_metrics_result()

            
            self.eval()

            
            if validation_data and self._should_eval(epoch, validation_freq):
                
                if getattr(self, "_eval_epoch_iterator", None) is None:
                    self._eval_epoch_iterator = EpochIterator(
                        x=val_x,
                        y=val_y,
                        sample_weight=val_sample_weight,
                        batch_size=validation_batch_size or batch_size,
                        steps_per_execution=self.steps_per_execution,
                        steps_per_epoch=validation_steps,
                        shuffle=False,
                    )
                val_logs = self.evaluate(
                    x=val_x,
                    y=val_y,
                    sample_weight=val_sample_weight,
                    batch_size=validation_batch_size or batch_size,
                    steps=validation_steps,
                    callbacks=callbacks,
                    return_dict=True,
                    _use_cached_eval_dataset=True,
                )
                val_logs = {
                    "val_" + name: val for name, val in val_logs.items()
                }
                epoch_logs.update(val_logs)

            callbacks.on_epoch_end(epoch, epoch_logs)
            training_logs = epoch_logs
            if self.stop_training:
                break

        if (
            isinstance(self.optimizer, optimizers_module.Optimizer)
            and epochs > 0
        ):
            self.optimizer.finalize_variable_values(self.trainable_weights)

        
        if getattr(self, "_eval_epoch_iterator", None) is not None:
            del self._eval_epoch_iterator
        callbacks.on_train_end(logs=training_logs)
        return self.history

    @traceback_utils.filter_traceback
    def evaluate(
        self,
        x=None,
        y=None,
        batch_size=None,
        verbose="auto",
        sample_weight=None,
        steps=None,
        callbacks=None,
        return_dict=False,
        **kwargs,
    ):
        
        use_cached_eval_dataset = kwargs.pop("_use_cached_eval_dataset", False)
        if kwargs:
            raise ValueError(f"Arguments not recognized: {kwargs}")

        if use_cached_eval_dataset:
            epoch_iterator = self._eval_epoch_iterator
        else:
            
            epoch_iterator = EpochIterator(
                x=x,
                y=y,
                sample_weight=sample_weight,
                batch_size=batch_size,
                steps_per_epoch=steps,
                shuffle=False,
                steps_per_execution=self.steps_per_execution,
            )

        if not all(layer.built for layer in self._flatten_layers()):
            
            for _, data in epoch_iterator.enumerate_epoch():
                data_batch = data[0]
                self._symbolic_build(data_batch)
                break

        
        if not isinstance(callbacks, callbacks_module.CallbackList):
            callbacks = callbacks_module.CallbackList(
                callbacks,
                add_history=True,
                add_progbar=verbose != 0,
                verbose=verbose,
                epochs=1,
                steps=epoch_iterator.num_batches,
                model=self,
            )

        
        self.eval()

        self.make_test_function()
        self.stop_evaluating = False
        callbacks.on_test_begin()
        logs = None
        self.reset_metrics()
        for step, data in epoch_iterator.enumerate_epoch():
            callbacks.on_test_batch_begin(step)
            logs = self.test_function(data)
            callbacks.on_test_batch_end(step, self._pythonify_logs(logs))
            if self.stop_evaluating:
                break
        logs = self.get_metrics_result()
        callbacks.on_test_end(logs)

        if return_dict:
            return logs
        return self._flatten_metrics_in_order(logs)

    @traceback_utils.filter_traceback
    def predict(
        self, x, batch_size=None, verbose="auto", steps=None, callbacks=None
    ):
        
        epoch_iterator = TorchEpochIterator(
            x=x,
            batch_size=batch_size,
            steps_per_epoch=steps,
            shuffle=False,
            steps_per_execution=self.steps_per_execution,
        )

        
        if not isinstance(callbacks, callbacks_module.CallbackList):
            callbacks = callbacks_module.CallbackList(
                callbacks,
                add_history=True,
                add_progbar=verbose != 0,
                verbose=verbose,
                epochs=1,
                steps=epoch_iterator.num_batches,
                model=self,
            )

        def append_to_outputs(batch_outputs, outputs):
            if outputs is None:
                outputs = tree.map_structure(
                    lambda batch_output: [batch_output],
                    batch_outputs,
                )
            else:
                tree.map_structure_up_to(
                    batch_outputs,
                    lambda output, batch_output: output.append(batch_output),
                    outputs,
                    batch_outputs,
                )
            return outputs

        
        self.eval()

        self.make_predict_function()
        self.stop_predicting = False
        callbacks.on_predict_begin()
        outputs = None
        for step, data in epoch_iterator.enumerate_epoch():
            callbacks.on_predict_batch_begin(step)
            batch_outputs = self.predict_function(data)
            outputs = append_to_outputs(batch_outputs, outputs)
            callbacks.on_predict_batch_end(step, {"outputs": batch_outputs})
            if self.stop_predicting:
                break
        callbacks.on_predict_end()
        outputs = tree.map_structure(backend.convert_to_numpy, outputs)
        return tree.map_structure_up_to(batch_outputs, np.concatenate, outputs)

    def train_on_batch(
        self,
        x,
        y=None,
        sample_weight=None,
        class_weight=None,
        return_dict=False,
    ):
        self._assert_compile_called("train_on_batch")
        if class_weight is not None:
            if sample_weight is not None:
                raise ValueError(
                    "Arguments `sample_weight` and `class_weight` "
                    "cannot be specified at the same time. "
                    f"Received: sample_weight={sample_weight}, "
                    f"class_weight={class_weight}"
                )
            sample_weight = data_adapter_utils.class_weight_to_sample_weights(
                y, class_weight
            )

        data = (x, y, sample_weight)

        
        self._symbolic_build(data)
        self.make_train_function()

        logs = self.train_function([data])
        logs = tree.map_structure(lambda x: np.array(x), logs)
        if return_dict:
            return logs
        return self._flatten_metrics_in_order(logs)

    def test_on_batch(
        self,
        x,
        y=None,
        sample_weight=None,
        return_dict=False,
    ):
        self._assert_compile_called("test_on_batch")

        data = (x, y, sample_weight)

        
        self._symbolic_build(data)
        self.make_test_function()

        logs = self.test_function([data])
        logs = tree.map_structure(lambda x: np.array(x), logs)
        if return_dict:
            return logs
        return self._flatten_metrics_in_order(logs)

    def predict_on_batch(self, x):
        self.make_predict_function()
        batch_outputs = self.predict_function([(x,)])
        batch_outputs = tree.map_structure(
            backend.convert_to_numpy, batch_outputs
        )
        return batch_outputs


class TorchEpochIterator(epoch_iterator.EpochIterator):
    def _get_iterator(self, return_type="auto"):
        if return_type == "auto" and isinstance(
            self.data_adapter, data_adapters.TorchDataLoaderAdapter
        ):
            return self.data_adapter.get_torch_dataloader()

        return super()._get_iterator(return_type)