add mnist_classifier/ (checkpoint + code)

mnist_classifier/classifier.py

@@ -0,0 +1,20 @@
+import jax
+from flax import linen as nn
+
+Array = jax.Array
+
+
+class Classifier(nn.Module):
+    num_classes: int
+    base_channels: int = 16
+
+    @nn.compact
+    def __call__(self, x: Array) -> Array:
+        h = nn.Conv(self.base_channels, kernel_size=(3, 3))(x)
+        h = nn.relu(h)
+        h = nn.max_pool(h, window_shape=(2, 2), strides=(2, 2))
+        h = nn.Conv(self.base_channels * 2, kernel_size=(3, 3))(h)
+        h = nn.relu(h)
+        h = nn.max_pool(h, window_shape=(2, 2), strides=(2, 2))
+        h = h.reshape(h.shape[0], -1)
+        return nn.Dense(self.num_classes)(h)

mnist_classifier/load.py

@@ -0,0 +1,13 @@
+import pickle
+
+from classifier import Classifier
+
+with open("mnist_classifier.pkl", "rb") as f:
+    bundle = pickle.load(f)
+meta = bundle["meta"]
+params = bundle["params"]
+
+model = Classifier(
+    num_classes=meta["num_classes"],
+    base_channels=meta["base_channels"],
+)

mnist_classifier/mnist_classifier.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2f68078916e0751319df3dd4280be0293bdbeff895a948b724cd9ded776846aa
+size 101781