model.export() results in ModuleNotFoundError: No module named 'tensorflow'

[tanwarsh]

Hi,

I am using Torch as the backend with Keras 3. I followed the guide at https://keras.io/guides/custom_train_step_in_torch/, but encountered the error "No module named 'tensorflow'" when trying to save the model after training.

Versions:

keras==3.8.0
torch==2.4.1
torchvision==0.19.1

Error can be reproduced by running the script provided below.

Is tensorflow still required just to save the model?

import os

os.environ["KERAS_BACKEND"] = "torch"

import torch
import keras
from keras import layers
import numpy as np


class CustomModel(keras.Model):
    def train_step(self, data):
        # Unpack the data. Its structure depends on your model and
        # on what you pass to `fit()`.
        if len(data) == 3:
            x, y, sample_weight = data
        else:
            sample_weight = None
            x, y = data

        # Call torch.nn.Module.zero_grad() to clear the leftover gradients
        # for the weights from the previous train step.
        self.zero_grad()

        # Compute loss
        y_pred = self(x, training=True)  # Forward pass
        loss = self.compute_loss(
            y=y,
            y_pred=y_pred,
            sample_weight=sample_weight,
        )

        # Call torch.Tensor.backward() on the loss to compute gradients
        # for the weights.
        loss.backward()

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

        # Update weights
        with torch.no_grad():
            self.optimizer.apply(gradients, trainable_weights)

        # 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, sample_weight=sample_weight)

        # Return a dict mapping metric names to current value
        # Note that it will include the loss (tracked in self.metrics).
        return {m.name: m.result() for m in self.metrics}

    def test_step(self, data):
        # Unpack the data
        x, y = data
        # Compute predictions
        y_pred = self(x, training=False)
        # Updates the metrics tracking the loss
        loss = self.compute_loss(y=y, y_pred=y_pred)
        # Update the metrics.
        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.
        # Note that it will include the loss (tracked in self.metrics).
        return {m.name: m.result() for m in self.metrics}

# Construct an instance of CustomModel
inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = CustomModel(inputs, outputs)

# We don't pass a loss or metrics here.
model.compile(loss="categorical_crossentropy",
                optimizer="adam",
                metrics=["accuracy"])

# Just use `fit` as usual -- you can use callbacks, etc.
x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.fit(x, y, epochs=5)

model.export("./")

Error:

Traceback (most recent call last):
  File "/home/perfuser/shailesh/torch_3feb/lib/python3.11/site-packages/keras/src/utils/module_utils.py", line 27, in initialize
    self.module = importlib.import_module(self.name)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.11/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<frozen importlib._bootstrap>", line 1204, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1176, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1140, in _find_and_load_unlocked
ModuleNotFoundError: No module named 'tensorflow'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/perfuser/shailesh/openfl_pytoch_latest_3_feb/torch_demo.py", line 162, in <module>
    model.export("./best.pbuf")
  File "/home/perfuser/shailesh/torch_3feb/lib/python3.11/site-packages/keras/src/models/model.py", line 539, in export
    export_saved_model(
  File "/home/perfuser/shailesh/torch_3feb/lib/python3.11/site-packages/keras/src/export/saved_model.py", line 624, in export_saved_model
    export_archive = ExportArchive()
                     ^^^^^^^^^^^^^^^
  File "/home/perfuser/shailesh/torch_3feb/lib/python3.11/site-packages/keras/src/export/saved_model.py", line 118, in __init__
    self.tensorflow_version = tf.__version__
                              ^^^^^^^^^^^^^^
  File "/home/perfuser/shailesh/torch_3feb/lib/python3.11/site-packages/keras/src/utils/module_utils.py", line 35, in __getattr__
    self.initialize()
  File "/home/perfuser/shailesh/torch_3feb/lib/python3.11/site-packages/keras/src/utils/module_utils.py", line 29, in initialize
    raise ImportError(self.import_error_msg)
ImportError: This requires the tensorflow module. You can install it via `pip install tensorflow`

if I use model.save instead of export, it result in error listed below.

model.save("./best.pbuf")

Traceback (most recent call last):
  File "/home/perfuser/shailesh/openfl_pytoch_latest_3_feb/torch_demo.py", line 162, in <module>
    model.save("./best.pbuf")
  File "/home/perfuser/shailesh/torch_3feb/lib/python3.11/site-packages/keras/src/utils/traceback_utils.py", line 122, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/perfuser/shailesh/torch_3feb/lib/python3.11/site-packages/keras/src/saving/saving_api.py", line 114, in save_model
    raise ValueError(
ValueError: Invalid filepath extension for saving. Please add either a `.keras` extension for the native Keras format (recommended) or a `.h5` extension. Use `model.export(filepath)` if you want to export a SavedModel for use with TFLite/TFServing/etc. Received: filepath=./best.pbuf.

pip freeze

absl-py==2.1.0
certifi==2025.1.31
cffi==1.17.1
charset-normalizer==3.4.1
click==8.1.8
cloudpickle==3.1.1
cryptography==44.0.0
dynaconf==3.2.7
filelock==3.17.0
flatten-json==0.1.14
fsspec==2025.2.0
grpcio==1.65.5
h5py==3.12.1
idna==3.10
Jinja2==3.1.5
joblib==1.4.2
keras==3.8.0
markdown-it-py==3.0.0
MarkupSafe==3.0.2
mdurl==0.1.2
ml_dtypes==0.5.1
mpmath==1.3.0
namex==0.0.8
networkx==3.4.2
numpy==2.2.2
nvidia-cublas-cu12==12.1.3.1
nvidia-cuda-cupti-cu12==12.1.105
nvidia-cuda-nvrtc-cu12==12.1.105
nvidia-cuda-runtime-cu12==12.1.105
nvidia-cudnn-cu12==9.1.0.70
nvidia-cufft-cu12==11.0.2.54
nvidia-curand-cu12==10.3.2.106
nvidia-cusolver-cu12==11.4.5.107
nvidia-cusparse-cu12==12.1.0.106
nvidia-nccl-cu12==2.20.5
nvidia-nvjitlink-cu12==12.8.61
nvidia-nvtx-cu12==12.1.105
openfl @ file:///home/perfuser/shailesh/openfl_pytoch_latest_3_feb
optree==0.14.0
packaging==24.2
pandas==2.2.3
pillow==11.1.0
protobuf==5.29.3
psutil==6.1.1
pycparser==2.22
Pygments==2.19.1
python-dateutil==2.9.0.post0
pytz==2025.1
PyYAML==6.0.2
requests==2.32.3
rich==13.9.4
scikit-learn==1.6.1
scipy==1.15.1
six==1.17.0
sympy==1.13.3
tensorboardX==2.6.2.2
threadpoolctl==3.5.0
torch==2.4.1
torchvision==0.19.1
tqdm==4.67.1
triton==3.0.0
typing_extensions==4.12.2
tzdata==2025.1
urllib3==2.3.0

[MasterSkepticista]

• model.export expects an argument format, which can be either tf_saved_model or onnx. TensorFlow is needed for the default option, which is tf_saved_model.
• model.save only supports .keras or .h5 file formats.

[tanwarsh]

thank you @MasterSkepticista