Optax Cifar 10 training (#103)

* cifar10 training

* optax

---------

Co-authored-by: samho <>

experiments/mnist/optax_cifar_training.py

@@ -0,0 +1,159 @@
+# Copyright 2018 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A basic MNIST example using Numpy and JAX.
+
+The primary aim here is simplicity and minimal dependencies.
+"""
+
+
+import time
+
+import datasets
+import jax
+import jax.numpy as jnp
+import numpy as np
+import numpy.random as npr
+import optax
+from jax import grad, jit, lax
+
+import jax_scaled_arithmetics as jsa
+
+
+def logsumexp(a, axis=None, keepdims=False):
+ dims = (axis,)
+ amax = jnp.max(a, axis=dims, keepdims=keepdims)
+ # FIXME: not proper scale propagation, introducing NaNs
+ # amax = lax.stop_gradient(lax.select(jnp.isfinite(amax), amax, lax.full_like(amax, 0)))
+ amax = lax.stop_gradient(amax)
+ out = lax.sub(a, amax)
+ out = lax.exp(out)
+ out = lax.add(lax.log(jnp.sum(out, axis=dims, keepdims=keepdims)), amax)
+ return out
+
+
+def init_random_params(scale, layer_sizes, rng=npr.RandomState(0)):
+ return [(scale * rng.randn(m, n), scale * rng.randn(n)) for m, n, in zip(layer_sizes[:-1], layer_sizes[1:])]
+
+
+def print_mean_std(name, v):
+ data, scale = jsa.lax.get_data_scale(v)
+ # Always use np.float32, to avoid floating errors in descaling + stats.
+ v = jsa.asarray(data, dtype=np.float32)
+ m, s = np.mean(v), np.std(v)
+ # print(data)
+ print(f"{name}: MEAN({m:.4f}) / STD({s:.4f}) / SCALE({scale:.4f})")
+
+
+def predict(params, inputs):
+ activations = inputs
+ for w, b in params[:-1]:
+ # Matmul + relu
+ outputs = jnp.dot(activations, w) + b
+ activations = jnp.maximum(outputs, 0)
+
+ final_w, final_b = params[-1]
+ logits = jnp.dot(activations, final_w) + final_b
+
+ # jsa.ops.debug_callback(partial(print_mean_std, "Logits"), logits)
+ # (logits,) = jsa.ops.debug_callback_grad(partial(print_mean_std, "LogitsGrad"), logits)
+
+ # Dynamic rescaling of the gradient, as logits gradient not properly scaled.
+ logits = jsa.ops.dynamic_rescale_l2_grad(logits)
+ output = logits - logsumexp(logits, axis=1, keepdims=True)
+
+ return output
+
+
+def loss(params, batch):
+ inputs, targets = batch
+ preds = predict(params, inputs)
+ return -jnp.mean(jnp.sum(preds * targets, axis=1))
+
+
+def accuracy(params, batch):
+ inputs, targets = batch
+ target_class = jnp.argmax(targets, axis=1)
+ predicted_class = jnp.argmax(predict(params, inputs), axis=1)
+ return jnp.mean(predicted_class == target_class)
+
+
+if __name__ == "__main__":
+ width = 256
+ lr = 1e-3
+ use_autoscale = False
+ training_dtype = np.float32
+ autoscale = jsa.autoscale if use_autoscale else lambda f: f
+
+ layer_sizes = [3072, width, width, 10]
+ param_scale = 1.0
+ num_epochs = 10
+ batch_size = 128
+ scale_dtype = np.float32
+
+ train_images, train_labels, test_images, test_labels = datasets.cifar()
+ num_train = train_images.shape[0]
+ num_complete_batches, leftover = divmod(num_train, batch_size)
+ num_batches = num_complete_batches + bool(leftover)
+ # num_batches = 2
+
+ def data_stream():
+ rng = npr.RandomState(0)
+ while True:
+ perm = rng.permutation(num_train)
+ for i in range(num_batches):
+ batch_idx = perm[i * batch_size : (i + 1) * batch_size]
+ yield train_images[batch_idx], train_labels[batch_idx]
+
+ batches = data_stream()
+ params = init_random_params(param_scale, layer_sizes)
+ params = jax.tree_map(lambda v: v.astype(training_dtype), params)
+ # Transform parameters to `ScaledArray` and proper dtype.
+ optimizer = optax.adam(learning_rate=lr, eps=1e-5)
+ opt_state = optimizer.init(params)
+
+ if use_autoscale:
+ params = jsa.as_scaled_array(params, scale=scale_dtype(param_scale))
+
+ params = jax.tree_map(lambda v: v.astype(training_dtype), params, is_leaf=jsa.core.is_scaled_leaf)
+
+ @jit
+ @autoscale
+ def update(params, batch, opt_state):
+ grads = grad(loss)(params, batch)
+ updates, opt_state = optimizer.update(grads, opt_state)
+ params = optax.apply_updates(params, updates)
+ return params, opt_state
+
+ for epoch in range(num_epochs):
+ start_time = time.time()
+ for _ in range(num_batches):
+ batch = next(batches)
+ # Scaled micro-batch + training dtype cast.
+ if use_autoscale:
+ batch = jsa.as_scaled_array(batch, scale=scale_dtype(param_scale))
+ batch = jax.tree_map(lambda v: v.astype(training_dtype), batch, is_leaf=jsa.core.is_scaled_leaf)
+
+ with jsa.AutoScaleConfig(rounding_mode=jsa.Pow2RoundMode.DOWN, scale_dtype=scale_dtype):
+ params, opt_state = update(params, batch, opt_state)
+
+ epoch_time = time.time() - start_time
+
+ # Evaluation in float32, for consistency.
+ raw_params = jsa.asarray(params, dtype=np.float32)
+ train_acc = accuracy(raw_params, (train_images, train_labels))
+ test_acc = accuracy(raw_params, (test_images, test_labels))
+ print(f"Epoch {epoch} in {epoch_time:0.2f} sec")
+ print(f"Training set accuracy {train_acc:0.5f}")
+ print(f"Test set accuracy {test_acc:0.5f}")

jax_scaled_arithmetics/lax/scaled_ops_common.py

@@ -316,3 +316,17 @@ def scaled_cos(val: ScaledArray) -> ScaledArray:
 @core.register_scaled_lax_op
 def scaled_sin(val: ScaledArray) -> ScaledArray:
  return scaled_op_default_translation(lax.sin_p, [val])
+
+
+@core.register_scaled_lax_op
+def scaled_integer_pow(A: ScaledArray, y: int) -> ScaledArray:
+ output_scale = lax.integer_pow(A.scale, y)
+ output_data = lax.integer_pow(A.data, y)
+ return ScaledArray(output_data, output_scale)
+
+
+@core.register_scaled_lax_op
+def scaled_sqrt(val: ScaledArray) -> ScaledArray:
+ output_scale = lax.sqrt(val.scale)
+ output_data = lax.sqrt(val.data)
+ return ScaledArray(output_data, output_scale)