Keras 3 API documentation
Models API Layers API The base Layer class Layer activations Layer weight initializers Layer weight regularizers Layer weight constraints Core layers Convolution layers Pooling layers Recurrent layers Preprocessing layers Normalization layers Regularization layers Attention layers Reshaping layers Merging layers Activation layers Backend-specific layers Callbacks API Ops API Optimizers Metrics Losses Data loading Built-in small datasets Keras Applications Mixed precision Multi-device distribution RNG API Rematerialization Utilities
Keras 2 API documentation
► Keras 3 API documentation / Layers API / Backend-specific layers / TorchModuleWrapper layer

TorchModuleWrapper layer

[source]

TorchModuleWrapper class

keras.layers.TorchModuleWrapper(module, name=None, output_shape=None, **kwargs)

Torch module wrapper layer.

TorchModuleWrapper is a wrapper class that can turn any torch.nn.Module into a Keras layer, in particular by making its parameters trackable by Keras.

TorchModuleWrapper is only compatible with the PyTorch backend and cannot be used with the TensorFlow or JAX backends.

Arguments

  • module: torch.nn.Module instance. If it's a LazyModule instance, then its parameters must be initialized before passing the instance to TorchModuleWrapper (e.g. by calling it once).
  • output_shape :The shape of the output of this layer. It helps Keras perform automatic shape inference.
  • name: The name of the layer (string).

Example

Here's an example of how the TorchModuleWrapper can be used with vanilla PyTorch modules.

import torch
import torch.nn as nn
import torch.nn.functional as F

import keras
from keras.layers import TorchModuleWrapper

class Classifier(keras.Model):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # Wrap `torch.nn.Module`s with `TorchModuleWrapper`
        # if they contain parameters
        self.conv1 = TorchModuleWrapper(
            nn.Conv2d(in_channels=1, out_channels=32, kernel_size=(3, 3))
        )
        self.conv2 = TorchModuleWrapper(
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=(3, 3))
        )
        self.pool = nn.MaxPool2d(kernel_size=(2, 2))
        self.flatten = nn.Flatten()
        self.dropout = nn.Dropout(p=0.5)
        self.fc = TorchModuleWrapper(nn.Linear(1600, 10))

    def call(self, inputs):
        x = F.relu(self.conv1(inputs))
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        x = self.flatten(x)
        x = self.dropout(x)
        x = self.fc(x)
        return F.softmax(x, dim=1)


model = Classifier()
model.build((1, 28, 28))
print("# Output shape", model(torch.ones(1, 1, 28, 28).to("cuda")).shape)

model.compile(
    loss="sparse_categorical_crossentropy",
    optimizer="adam",
    metrics=["accuracy"]
)
model.fit(train_loader, epochs=5)
TorchModuleWrapper layer
TorchModuleWrapper class