small fix

PiperOrigin-RevId: 691467279

Author: xingyousong

optformer/embed_then_regress/checkpointing.py

@@ -21,18 +21,16 @@
 
 
 def get_checkpoint_manager(
-    workdir: epath.PathLike,
+    workdir: epath.PathLike, **options_kwargs
 ) -> orbax_checkpoint.CheckpointManager:
   """Sets up Orbax checkpointing."""
-  # The keys in this dict should match the keys in `checkpointed_state`.
-  checkpointers = dict(
-      train_state=orbax_checkpoint.PyTreeCheckpointer(),
-  )
   checkpoint_dir = epath.Path(workdir) / 'checkpoints'
   return orbax_checkpoint.CheckpointManager(
       checkpoint_dir,
-      checkpointers=checkpointers,
-      options=orbax_checkpoint.CheckpointManagerOptions(create=True),
+      checkpointers={'train_state': orbax_checkpoint.PyTreeCheckpointer()},
+      options=orbax_checkpoint.CheckpointManagerOptions(
+          create=True, **options_kwargs
+      ),
   )
 
 

optformer/embed_then_regress/configs.py

@@ -20,6 +20,7 @@
 from typing import Callable
 from flax import linen as nn
 import jax
+import jax.numpy as jnp
 import jaxtyping as jt
 import optax
 from optformer.embed_then_regress import icl_transformer
@@ -84,12 +85,11 @@ def create_model(
   ) -> icl_transformer.ICLTransformer:
 
     kwargs = dataclasses.asdict(self)
-    embedder_factory = embedder_config.create_embedder_factory()
-    std_transform_fn = self.create_std_transform_fn()
+    kwargs.pop('std_transform')
 
     return icl_transformer.ICLTransformer(
-        std_transform_fn=std_transform_fn,
-        embedder_factory=embedder_factory,
+        std_transform_fn=self.create_std_transform_fn(),
+        embedder_factory=embedder_config.create_embedder_factory(),
         **kwargs,
     )
 
@@ -98,17 +98,17 @@ def create_std_transform_fn(
   ) -> Callable[[jt.Float[jax.Array, '*A']], jt.Float[jax.Array, '*A']]:
     """Creates std transform function."""
     if self.std_transform == 'exp':
-      return jax.numpy.exp
+      return jnp.exp
     elif self.std_transform == 'exp10':
-      return lambda x: jax.numpy.exp(10.0 * x)
+      return lambda x: jnp.exp(10.0 * x)
     elif self.std_transform == 'softplus':
       return jax.nn.softplus
     elif self.std_transform == 'softplus10':
       return lambda x: jax.nn.softplus(10.0 * x)
     elif self.std_transform == 'abs':
-      return jax.numpy.abs
+      return jnp.abs
     elif self.std_transform == 'abs10':
-      return lambda x: jax.numpy.abs(10.0 * x)
+      return lambda x: jnp.abs(10.0 * x)
     elif self.std_transform == 'shifted_relu':
       return lambda x: jax.nn.relu(x + 1.0)
     elif self.std_transform == 'shifted_relu10':
@@ -131,7 +131,7 @@ class TrainingConfig:
   seed: int = 42
 
   validation_interval: int = 100
-  checkpoint_interval: int = 100
+  max_to_keep_ckpts: int = 5
   workdir = '../checkpoints'
 
   def create_optimizer(self) -> optax.GradientTransformation:
@@ -186,7 +186,6 @@ def wrap_ds(
   ) -> tf.data.Dataset:
     """This should be used at the trainer level."""
     ds = self._tokenize_ds(ds)
-    ds = ds.shard(jax.process_count(), jax.process_index())
     ds = ds.repeat()
     ds = ds.shuffle(buffer_size=self.buffer_size)
 

optformer/embed_then_regress/icl_transformer.py

@@ -169,7 +169,7 @@ def __call__(
       out = layer(out, mask, deterministic, rng)
 
     mean, std = jnp.split(self.mean_logstd_head(out), 2, axis=-1)  # [B L 1]
-    std = self.std_transform_fn(self.std_transform)(std) + EPS
+    std = self.std_transform_fn(std) + EPS
 
     mean = jnp.squeeze(mean, axis=-1)
     std = jnp.squeeze(std, axis=-1)

optformer/embed_then_regress/train.py

@@ -155,7 +155,8 @@ def aggregate_metrics(
   """Aggregates metrics (possibly from multiple gradient accumulation steps)."""
   if isinstance(metrics, list):
     metrics = jax.tree.map(lambda *args: jnp.stack(args), *metrics)
-  return jax.tree.map(jnp.mean, metrics)
+  metrics = jax.tree.map(jnp.mean, metrics)
+  return {k: float(v) for k, v in metrics.items()}
 
 
 def train(
@@ -194,7 +195,12 @@ def train(
   )
 
   # Set up checkpointing
-  checkpoint_manager = ckpt_lib.get_checkpoint_manager(train_config.workdir)
+  checkpoint_manager = ckpt_lib.get_checkpoint_manager(
+      train_config.workdir,
+      max_to_keep=train_config.max_to_keep_ckpts,
+      best_fn=lambda metrics: metrics['eval_loss'],
+      best_mode='min',
+  )
   # Restore if available.
   train_state = ckpt_lib.restore_train_state(
       train_config.workdir, init_train_state
@@ -206,20 +212,20 @@ def train(
   eff_step = int(unreplicate(train_state.step)) // grad_accum_steps
 
   while eff_step < train_config.max_steps:
-    if eff_step % train_config.checkpoint_interval == 0:
+    if eff_step % train_config.validation_interval == 0:
+      valid_agg_metrics = aggregate_metrics([
+          p_eval_step(train_state, next(valid_it))
+          for _ in range(grad_accum_steps)
+      ])
+      writer.write_scalars(eff_step, valid_agg_metrics)
+
       ckpt_train_state = unreplicate(train_state)
       checkpoint_manager.save(
           eff_step,
           items=dict(train_state=jax.tree.map(np.array, ckpt_train_state)),
+          metrics=valid_agg_metrics,
       )
 
-    if eff_step % train_config.validation_interval == 0:
-      all_valid_metrics = [
-          p_eval_step(train_state, next(valid_it))
-          for _ in range(grad_accum_steps)
-      ]
-      writer.write_scalars(eff_step, aggregate_metrics(all_valid_metrics))
-
     all_train_metrics = []
     for _ in range(grad_accum_steps):
       train_state, train_metrics = p_train_step(train_state, next(train_it))