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 args 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 name the name of the layer string examples here s an example of how the torchmodulewrapper can be used with vanilla pytorch modules python import torch nn as nn import torch nn functional as f import keras from keras layers import torchmodulewrapper class classifierkeras model def initself kwargs super initkwargs wrap torch nn modules with torchmodulewrapper if they contain parameters self conv1 torchmodulewrapper nn conv2dinchannels1 outchannels32 kernelsize3 3 self conv2 torchmodulewrapper nn conv2dinchannels32 outchannels64 kernelsize3 3 self pool nn maxpool2dkernelsize2 2 self flatten nn flatten self dropout nn dropoutp0 5 self fc torchmodulewrappernn linear1600 10 def callself inputs x f reluself conv1inputs x self poolx x f reluself conv2x x self poolx x self flattenx x self dropoutx x self fcx return f softmaxx dim1 model classifier model build1 28 28 printoutput shape modeltorch ones1 1 28 28 tocuda shape model compile losssparsecategoricalcrossentropy optimizeradam metricsaccuracy model fittrainloader epochs5 the variable will reuse the raw param and simply wrap it saves model s state from statedict model parameters excludes some of model s state like batchnorm mean and variance so use statedict to obtain all of model s state loads model s state via statedict statedict for key in store keys if isinstancekey bytes key key decode statedictkey converttotensorstorekey self module loadstatedictstatedict def getconfigself baseconfig super getconfig import torch buffer io bytesio torch saveself module buffer config module buffer getvalue return baseconfig config classmethod def fromconfigcls config import torch if module in config buffer io bytesioconfigmodule configmodule torch loadbuffer return clsconfig def nogradorigfunc import torch if parsetorch version parse2 1 0 return torch nogradorigfunc else return origfunc 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 args 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 name the name of the layer string examples here s an example of how the torchmodulewrapper can be used with vanilla pytorch modules python 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 the variable will reuse the raw param and simply wrap it saves model s state from state_dict model parameters excludes some of model s state like batchnorm mean and variance so use state_dict to obtain all of model s state loads model s state via state_dict
import io

from packaging.version import parse

from keras.api_export import keras_export
from keras.layers import Layer
from keras.ops import convert_to_numpy
from keras.ops import convert_to_tensor


@keras_export("keras.layers.TorchModuleWrapper")
class TorchModuleWrapper(Layer):
    

    def __init__(self, module, name=None, **kwargs):
        super().__init__(name=name, **kwargs)
        import torch.nn as nn

        from keras.backend.torch.core import get_device

        if (
            isinstance(module, nn.modules.lazy.LazyModuleMixin)
            and module.has_uninitialized_params()
        ):
            raise ValueError(
                "LazyModules are not supported unless they "
                "are already initialized. "
                f"Received uninitialized LazyModule: module={module}"
            )

        self.module = module.to(get_device())
        self._track_module_parameters()

    def parameters(self, recurse=True):
        return self.module.parameters(recurse=recurse)

    def _track_module_parameters(self):
        from keras.backend.torch import Variable

        for param in self.module.parameters():
            
            
            variable = Variable(
                initializer=param, trainable=param.requires_grad
            )
            self._track_variable(variable)
        self.built = True

    def call(self, *args, **kwargs):
        return self.module.forward(*args, **kwargs)

    def save_own_variables(self, store):
        
        state_dict = self.module.state_dict()
        for key in state_dict.keys():
            store[key] = convert_to_numpy(state_dict[key])

    def load_own_variables(self, store):
        
        state_dict = {}
        for key in store.keys():
            if isinstance(key, bytes):
                key = key.decode()
            state_dict[key] = convert_to_tensor(store[key])
        self.module.load_state_dict(state_dict)

    def get_config(self):
        base_config = super().get_config()
        import torch

        buffer = io.BytesIO()
        torch.save(self.module, buffer)
        config = {"module": buffer.getvalue()}
        return {**base_config, **config}

    @classmethod
    def from_config(cls, config):
        import torch

        if "module" in config:
            buffer = io.BytesIO(config["module"])
            config["module"] = torch.load(buffer)
        return cls(**config)


def no_grad(orig_func):
    import torch

    if parse(torch.__version__) >= parse("2.1.0"):
        return torch.no_grad(orig_func)
    else:
        return orig_func