updated cifar10 example, mixer model typing

Author: homerjed

data/cifar10.py

@@ -57,15 +57,15 @@ def cifar10(path: str, key: Key, *, in_memory: bool = True) -> ScalerDataset:
         target_transform=transforms.Lambda(lambda x: x.float())
     )
 
-
     if in_memory:
         Xt, At = convert_torch_to_in_memory(train_dataset) 
         Xv, Av = convert_torch_to_in_memory(valid_dataset) 
 
         At = At.astype(jnp.float32)
         Av = Av.astype(jnp.float32)
 
-        process_fn = Scaler(x_min=Xt.min(), x_max=Xt.max())
+        # process_fn = Scaler(x_min=Xt.min(), x_max=Xt.max())
+        process_fn = Normer(x_mean=Xt.mean(), x_std=Xt.std())
 
         train_dataloader = InMemoryDataLoader(
             X=Xt, A=At, process_fn=process_fn, key=key_train) 

examples/cifar10.ipynb

@@ -466,7 +466,7 @@ def train_from_config(
                 # Plot losses etc
                 plot_metrics(train_losses, valid_losses, step, exp_dir)
 
-    return model
+    return ema_model if config.use_ema else model
 
 
 def train(
@@ -726,4 +726,4 @@ def train(
                 # Plot losses etc
                 plot_metrics(train_losses, valid_losses, step, exp_dir)
 
-    return model
+    return ema_model if use_ema else model

sbgm/_train.py

@@ -1,10 +1,11 @@
-from typing import Sequence, Optional, Union
+from typing import Sequence, Optional, Callable
 import jax
 import jax.numpy as jnp
 import jax.random as jr
 import einops
 import equinox as eqx
-from jaxtyping import Key, Array
+from jaxtyping import Key, Array, Float, jaxtyped
+from beartype import beartype as typechecker
 
 
 class AdaLayerNorm(eqx.Module):
@@ -134,7 +135,11 @@ class Mixer2d(eqx.Module):
     t1: float
     embedding_dim: int
     final_activation: callable
+    img_size: Sequence[int]
+    q_dim: int
+    a_dim: int
 
+    @jaxtyped(typechecker=typechecker)
     def __init__(
         self,
         img_size: Sequence[int],
@@ -145,11 +150,11 @@ def __init__(
         num_blocks: int,
         t1: float,
         embedding_dim: int = 8,
-        final_activation: Optional[Union[callable, str]] = None,
+        final_activation: Optional[Callable | str] = None,
         q_dim: Optional[int] = None,
         a_dim: Optional[int] = None,
         *,
-        key: Key
+        key: Key[jnp.ndarray, "..."]
     ):
         """
             A 2D MLP Mixer model.
@@ -207,6 +212,10 @@ def __init__(
         _input_size = input_size + q_dim if q_dim is not None else input_size
         _context_dim = embedding_dim + a_dim if a_dim is not None else embedding_dim
 
+        self.img_size = img_size
+        self.q_dim = q_dim
+        self.a_dim = a_dim
+
         self.conv_in = eqx.nn.Conv2d(
             _input_size, 
             hidden_size, 
@@ -237,15 +246,16 @@ def __init__(
         self.embedding_dim = embedding_dim
         self.final_activation = get_activation_fn(final_activation)
 
+    @jaxtyped(typechecker=typechecker)
     def __call__(
         self, 
-        t: Union[float, Array], 
-        y: Array, 
-        q: Optional[Array] = None, 
-        a: Optional[Array] = None, 
+        t: float | Float[Array, ""], 
+        y: Float[Array, "..."], 
+        q: Optional[Float[Array, "{self.q_dim} ..."]] = None, 
+        a: Optional[Float[Array, "{self.a_dim}"]] = None, 
         *, 
-        key: Optional[Key] = None
-    ) -> Array:
+        key: Optional[Key[jnp.ndarray, "..."]] = None
+    ) -> Float[Array, "..."]:
         _, height, width = y.shape
         t = jnp.atleast_1d(t / self.t1)
         t = get_timestep_embedding(t, embedding_dim=self.embedding_dim)