Add (de)serialization code for TimeSeriesTransformer

Author: vpratz

bayesflow/networks/transformers/time_series_transformer.py

@@ -1,8 +1,8 @@
 import keras
 
 from bayesflow.types import Tensor
-from bayesflow.utils import check_lengths_same
-from bayesflow.utils.serialization import serializable
+from bayesflow.utils import check_lengths_same, model_kwargs
+from bayesflow.utils.serialization import deserialize, serializable, serialize
 
 from ..embeddings import Time2Vec, RecurrentEmbedding
 from ..summary_network import SummaryNetwork
@@ -103,9 +103,22 @@ def __init__(
         # Pooling will be applied as a final step to the abstract representations obtained from set attention
         self.pooling = keras.layers.GlobalAvgPool1D()
         self.output_projector = keras.layers.Dense(summary_dim)
-        self.summary_dim = summary_dim
 
+        # store variables for serialization
+        self.summary_dim = summary_dim
+        self.embed_dims = embed_dims
+        self.num_heads = num_heads
+        self.mlp_depths = mlp_depths
+        self.mlp_widths = mlp_widths
+        self.dropout = dropout
+        self.mlp_activation = mlp_activation
+        self.kernel_initializer = kernel_initializer
+        self.use_bias = use_bias
+        self.layer_norm = layer_norm
+        self._time_embedding_arg = time_embedding
+        self.time_embed_dim = time_embed_dim
         self.time_axis = time_axis
+        self._kwargs = kwargs
 
     def call(self, input_sequence: Tensor, training: bool = False, **kwargs) -> Tensor:
         """Compresses the input sequence into a summary vector of size `summary_dim`.
@@ -147,3 +160,30 @@ def call(self, input_sequence: Tensor, training: bool = False, **kwargs) -> Tens
         summary = self.pooling(inp)
         summary = self.output_projector(summary)
         return summary
+
+    @classmethod
+    def from_config(cls, config, custom_objects=None):
+        return cls(**deserialize(config, custom_objects=custom_objects))
+
+    def get_config(self):
+        base_config = super().get_config()
+        base_config = model_kwargs(base_config)
+
+        config = {
+            "summary_dim": self.summary_dim,
+            "embed_dims": self.embed_dims,
+            "num_heads": self.num_heads,
+            "mlp_depths": self.mlp_depths,
+            "mlp_widths": self.mlp_widths,
+            "dropout": self.dropout,
+            "mlp_activation": self.mlp_activation,
+            "kernel_initializer": self.kernel_initializer,
+            "use_bias": self.use_bias,
+            "layer_norm": self.layer_norm,
+            "time_embedding": self._time_embedding_arg,
+            "time_embed_dim": self.time_embed_dim,
+            "time_axis": self.time_axis,
+            **self._kwargs,
+        }
+
+        return base_config | serialize(config)