fix: make NNs pickleable

src/safeds/ml/nn/_internal_model.py

@@ -0,0 +1,60 @@
+# The class must not be nested inside a function, since pickle cannot serialize local classes. Because of this, the
+# slow import of torch must be on the global level. To still evaluate the torch import lazily, the class is moved to a
+# separate file.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from torch import Tensor, nn  # slow import
+
+from safeds._config import _init_default_device
+from safeds.ml.nn.converters._input_converter_image import _InputConversionImage
+from safeds.ml.nn.layers import FlattenLayer, Layer
+from safeds.ml.nn.layers._pooling2d_layer import _Pooling2DLayer
+
+if TYPE_CHECKING:
+    from safeds.ml.nn.converters import InputConversion
+    from safeds.ml.nn.typing import ModelImageSize
+
+
+# Use torch.compile once the following issues are resolved:
+# - https://github.com/pytorch/pytorch/issues/120233 (Python 3.12 support)
+# - https://github.com/triton-lang/triton/issues/1640 (Windows support)
+class _InternalModel(nn.Module):
+    def __init__(self, input_conversion: InputConversion, layers: list[Layer], is_for_classification: bool) -> None:
+        super().__init__()
+
+        _init_default_device()
+
+        self._layer_list = layers
+        internal_layers = []
+        previous_output_size = input_conversion._data_size
+
+        for layer in layers:
+            if previous_output_size is not None:
+                layer._set_input_size(previous_output_size)
+            elif isinstance(input_conversion, _InputConversionImage):
+                layer._set_input_size(input_conversion._data_size)
+            if isinstance(layer, FlattenLayer | _Pooling2DLayer):
+                internal_layers.append(layer._get_internal_layer())
+            else:
+                internal_layers.append(layer._get_internal_layer(activation_function="relu"))
+            previous_output_size = layer.output_size
+
+        if is_for_classification:
+            internal_layers.pop()
+            if isinstance(layers[-1].output_size, int) and layers[-1].output_size > 2:
+                internal_layers.append(layers[-1]._get_internal_layer(activation_function="none"))
+            else:
+                internal_layers.append(layers[-1]._get_internal_layer(activation_function="sigmoid"))
+        self._pytorch_layers = nn.Sequential(*internal_layers)
+
+    @property
+    def input_size(self) -> int | ModelImageSize:
+        return self._layer_list[0].input_size
+
+    def forward(self, x: Tensor) -> Tensor:
+        for layer in self._pytorch_layers:
+            x = layer(x)
+        return x

src/safeds/ml/nn/_model.py

@@ -32,7 +32,7 @@
 if TYPE_CHECKING:
     from collections.abc import Callable
 
-    from torch import Tensor, nn
+    from torch import nn
     from torch.nn import Module
     from transformers.image_processing_utils import BaseImageProcessor
 
@@ -206,6 +206,8 @@ def fit(
         import torch
         from torch import nn
 
+        from ._internal_model import _InternalModel  # Slow import on global level
+
         _init_default_device()
 
         if not self._input_conversion._is_fit_data_valid(train_data):
@@ -216,7 +218,7 @@ def fit(
 
         copied_model = copy.deepcopy(self)
         # TODO: How is this supposed to work with pre-trained models? Should the old weights be kept or discarded?
-        copied_model._model = _create_internal_model(self._input_conversion, self._layers, is_for_classification=False)
+        copied_model._model = _InternalModel(self._input_conversion, self._layers, is_for_classification=False)
         copied_model._input_size = copied_model._model.input_size
         copied_model._batch_size = batch_size
 
@@ -501,6 +503,8 @@ def fit(
         import torch
         from torch import nn
 
+        from ._internal_model import _InternalModel  # Slow import on global level
+
         _init_default_device()
 
         if not self._input_conversion._is_fit_data_valid(train_data):
@@ -511,7 +515,7 @@ def fit(
 
         copied_model = copy.deepcopy(self)
         # TODO: How is this supposed to work with pre-trained models? Should the old weights be kept or discarded?
-        copied_model._model = _create_internal_model(self._input_conversion, self._layers, is_for_classification=True)
+        copied_model._model = _InternalModel(self._input_conversion, self._layers, is_for_classification=True)
         copied_model._batch_size = batch_size
         copied_model._input_size = copied_model._model.input_size
 
@@ -617,53 +621,3 @@ def input_size(self) -> int | ModelImageSize | None:
         """The input size of the model."""
         # TODO: raise if not fitted, don't return None
         return self._input_size
-
-
-def _create_internal_model(
-    input_conversion: InputConversion[IFT, IPT],
-    layers: list[Layer],
-    is_for_classification: bool,
-) -> nn.Module:
-    from torch import nn
-
-    _init_default_device()
-
-    class _InternalModel(nn.Module):
-        def __init__(self, layers: list[Layer], is_for_classification: bool) -> None:
-            super().__init__()
-            self._layer_list = layers
-            internal_layers = []
-            previous_output_size = input_conversion._data_size
-
-            for layer in layers:
-                if previous_output_size is not None:
-                    layer._set_input_size(previous_output_size)
-                elif isinstance(input_conversion, _InputConversionImage):
-                    layer._set_input_size(input_conversion._data_size)
-                if isinstance(layer, FlattenLayer | _Pooling2DLayer):
-                    internal_layers.append(layer._get_internal_layer())
-                else:
-                    internal_layers.append(layer._get_internal_layer(activation_function="relu"))
-                previous_output_size = layer.output_size
-
-            if is_for_classification:
-                internal_layers.pop()
-                if isinstance(layers[-1].output_size, int) and layers[-1].output_size > 2:
-                    internal_layers.append(layers[-1]._get_internal_layer(activation_function="none"))
-                else:
-                    internal_layers.append(layers[-1]._get_internal_layer(activation_function="sigmoid"))
-            self._pytorch_layers = nn.Sequential(*internal_layers)
-
-        @property
-        def input_size(self) -> int | ModelImageSize:
-            return self._layer_list[0].input_size
-
-        def forward(self, x: Tensor) -> Tensor:
-            for layer in self._pytorch_layers:
-                x = layer(x)
-            return x
-
-    # Use torch.compile once the following issues are resolved:
-    # - https://github.com/pytorch/pytorch/issues/120233 (Python 3.12 support)
-    # - https://github.com/triton-lang/triton/issues/1640 (Windows support)
-    return _InternalModel(layers, is_for_classification)

src/safeds/ml/nn/layers/_convolutional2d_layer.py

@@ -4,82 +4,14 @@
 import sys
 from typing import TYPE_CHECKING, Any, Literal
 
-from safeds._config import _init_default_device
 from safeds._utils import _structural_hash
 
 from ._layer import Layer
 
 if TYPE_CHECKING:
-    from torch import Tensor, nn
-
-    from safeds.ml.nn.typing import ModelImageSize
-
-
-def _create_internal_model(
-    input_size: int,
-    output_size: int,
-    kernel_size: int,
-    activation_function: Literal["sigmoid", "relu", "softmax"],
-    padding: int,
-    stride: int,
-    transpose: bool,
-    output_padding: int = 0,
-) -> nn.Module:
     from torch import nn
 
-    _init_default_device()
-
-    class _InternalLayer(nn.Module):
-        def __init__(
-            self,
-            input_size: int,
-            output_size: int,
-            kernel_size: int,
-            activation_function: Literal["sigmoid", "relu", "softmax"],
-            padding: int,
-            stride: int,
-            transpose: bool,
-            output_padding: int,
-        ):
-            super().__init__()
-            if transpose:
-                self._layer = nn.ConvTranspose2d(
-                    in_channels=input_size,
-                    out_channels=output_size,
-                    kernel_size=kernel_size,
-                    padding=padding,
-                    stride=stride,
-                    output_padding=output_padding,
-                )
-            else:
-                self._layer = nn.Conv2d(
-                    in_channels=input_size,
-                    out_channels=output_size,
-                    kernel_size=kernel_size,
-                    padding=padding,
-                    stride=stride,
-                )
-            match activation_function:
-                case "sigmoid":
-                    self._fn = nn.Sigmoid()
-                case "relu":
-                    self._fn = nn.ReLU()
-                case "softmax":
-                    self._fn = nn.Softmax()
-
-        def forward(self, x: Tensor) -> Tensor:
-            return self._fn(self._layer(x))
-
-    return _InternalLayer(
-        input_size,
-        output_size,
-        kernel_size,
-        activation_function,
-        padding,
-        stride,
-        transpose,
-        output_padding,
-    )
+    from safeds.ml.nn.typing import ModelImageSize
 
 
 class Convolutional2DLayer(Layer):
@@ -107,6 +39,8 @@ def __init__(self, output_channel: int, kernel_size: int, *, stride: int = 1, pa
         self._output_size: ModelImageSize | None = None
 
     def _get_internal_layer(self, **kwargs: Any) -> nn.Module:
+        from ._internal_layers import _InternalConvolutional2DLayer  # slow import on global level
+
         if self._input_size is None:
             raise ValueError(
                 "The input_size is not yet set. The internal layer can only be created when the input_size is set.",
@@ -121,14 +55,15 @@ def _get_internal_layer(self, **kwargs: Any) -> nn.Module:
             )
         else:
             activation_function: Literal["sigmoid", "relu", "softmax"] = kwargs["activation_function"]
-        return _create_internal_model(
+        return _InternalConvolutional2DLayer(
             self._input_size.channel,
             self._output_channel,
             self._kernel_size,
             activation_function,
             self._padding,
             self._stride,
             transpose=False,
+            output_padding=0,
         )
 
     @property
@@ -254,6 +189,8 @@ def __init__(
         self._output_padding = output_padding
 
     def _get_internal_layer(self, **kwargs: Any) -> nn.Module:
+        from ._internal_layers import _InternalConvolutional2DLayer  # slow import on global level
+
         if self._input_size is None:
             raise ValueError(
                 "The input_size is not yet set. The internal layer can only be created when the input_size is set.",
@@ -268,7 +205,7 @@ def _get_internal_layer(self, **kwargs: Any) -> nn.Module:
             )
         else:
             activation_function: Literal["sigmoid", "relu", "softmax"] = kwargs["activation_function"]
-        return _create_internal_model(
+        return _InternalConvolutional2DLayer(
             self._input_size.channel,
             self._output_channel,
             self._kernel_size,

src/safeds/ml/nn/layers/_flatten_layer.py

@@ -3,32 +3,15 @@
 import sys
 from typing import TYPE_CHECKING, Any
 
-from safeds._config import _init_default_device
 from safeds._utils import _structural_hash
 from safeds.ml.nn.typing import ConstantImageSize
 
 from ._layer import Layer
 
 if TYPE_CHECKING:
-    from torch import Tensor, nn
-
-    from safeds.ml.nn.typing import ModelImageSize
-
-
-def _create_internal_model() -> nn.Module:
     from torch import nn
 
-    _init_default_device()
-
-    class _InternalLayer(nn.Module):
-        def __init__(self) -> None:
-            super().__init__()
-            self._layer = nn.Flatten()
-
-        def forward(self, x: Tensor) -> Tensor:
-            return self._layer(x)
-
-    return _InternalLayer()
+    from safeds.ml.nn.typing import ModelImageSize
 
 
 class FlattenLayer(Layer):
@@ -39,7 +22,9 @@ def __init__(self) -> None:
         self._output_size: int | None = None
 
     def _get_internal_layer(self, **kwargs: Any) -> nn.Module:  # noqa: ARG002
-        return _create_internal_model()
+        from ._internal_layers import _InternalFlattenLayer  # Slow import on global level
+
+        return _InternalFlattenLayer()
 
     @property
     def input_size(self) -> ModelImageSize:

src/safeds/ml/nn/layers/_forward_layer.py

@@ -2,15 +2,14 @@
 
 from typing import TYPE_CHECKING, Any
 
-from safeds._config import _init_default_device
 from safeds._utils import _structural_hash
 from safeds._validation import _check_bounds, _ClosedBound
 from safeds.ml.nn.typing import ModelImageSize
 
 from ._layer import Layer
 
 if TYPE_CHECKING:
-    from torch import Tensor, nn
+    from torch import nn
 
 
 class ForwardLayer(Layer):
@@ -36,6 +35,8 @@ def __init__(self, neuron_count: int):
         self._output_size = neuron_count
 
     def _get_internal_layer(self, **kwargs: Any) -> nn.Module:
+        from ._internal_layers import _InternalForwardLayer  # Slow import on global level
+
         if "activation_function" not in kwargs:
             raise ValueError(
                 "The activation_function is not set. The internal layer can only be created when the activation_function is provided in the kwargs.",
@@ -46,7 +47,7 @@ def _get_internal_layer(self, **kwargs: Any) -> nn.Module:
         if self._input_size is None:
             raise ValueError("The input_size is not yet set.")
 
-        return _create_internal_model(self._input_size, self._output_size, activation_function)
+        return _InternalForwardLayer(self._input_size, self._output_size, activation_function)
 
     @property
     def input_size(self) -> int:
@@ -95,30 +96,3 @@ def __sizeof__(self) -> int:
         import sys
 
         return sys.getsizeof(self._input_size) + sys.getsizeof(self._output_size)
-
-
-def _create_internal_model(input_size: int, output_size: int, activation_function: str) -> nn.Module:
-    from torch import nn
-
-    _init_default_device()
-
-    class _InternalLayer(nn.Module):
-        def __init__(self, input_size: int, output_size: int, activation_function: str):
-            super().__init__()
-            self._layer = nn.Linear(input_size, output_size)
-            match activation_function:
-                case "sigmoid":
-                    self._fn = nn.Sigmoid()
-                case "relu":
-                    self._fn = nn.ReLU()
-                case "softmax":
-                    self._fn = nn.Softmax()
-                case "none":
-                    self._fn = None
-                case _:
-                    raise ValueError("Unknown Activation Function: " + activation_function)
-
-        def forward(self, x: Tensor) -> Tensor:
-            return self._fn(self._layer(x)) if self._fn is not None else self._layer(x)
-
-    return _InternalLayer(input_size, output_size, activation_function)