Adding logging for the number of parameters and optimizer state.

PiperOrigin-RevId: 533111618

Author: botev

examples/training.py

@@ -126,6 +126,9 @@ def __init__(
 
     self._params, self._state, self._opt_state = None, None, None
     self._python_step = 0
+    self._num_tensors = 0
+    self._num_parameters = 0
+    self._optimizer_state_size = 0
 
   def log_machines_setup(self):
     """Logs the machine setup for the experiment."""
@@ -377,12 +380,14 @@ def maybe_initialize_state(self):
     init_rng = kfac_jax.utils.replicate_all_local_devices(self.init_rng)
     # Initialize parameters and optional state
     params_rng, optimizer_rng = kfac_jax.utils.p_split(init_rng)
+    logging.info("Initializing parameters.")
     if self.has_func_state:
       self._params, self._state = self.params_init(params_rng, self.init_batch)
     else:
       self._params = self.params_init(params_rng, self.init_batch)
 
     # Initialize optimizer state
+    logging.info("Initializing optimizer state.")
     self._opt_state = self.optimizer.init(
         self._params, optimizer_rng, self.init_batch, self._state
     )
@@ -391,6 +396,42 @@ def maybe_initialize_state(self):
       # Needed for checkpointing
       self._state = ()
 
+    # Log parameters
+    self._num_tensors = 0
+    self._num_parameters = 0
+    logging.info("%s %s %s", "=" * 20, "Parameters", "=" * 20)
+    for path, var in jax.tree_util.tree_flatten_with_path(self._params)[0]:
+      # Because of pmap
+      var = var[0]
+      logging.info(
+          "%s - %s, %s",
+          "-".join(str(p)[2:-2] for p in path),
+          var.shape,
+          var.dtype,
+      )
+      self._num_parameters = self._num_parameters + var.size
+      self._num_tensors = self._num_tensors + 1
+    logging.info("Total parameters: %s", f"{self._num_parameters:,}")
+
+    # Log optimizer state
+    self._optimizer_state_size = 0
+    logging.info("%s %s %s", "=" * 20, "Optimizer State", "=" * 20)
+    easy_state = kfac_jax.utils.serialize_state_tree(self._opt_state)
+    for path, var in jax.tree_util.tree_flatten_with_path(easy_state)[0]:
+      if isinstance(var, str):
+        # For __class__ entries
+        continue
+      # Because of pmap
+      var = var[0]
+      logging.info(
+          "%s - %s, %s",
+          "/".join(str(p)[2:-2] for p in path),
+          var.shape,
+          var.dtype,
+      )
+      self._optimizer_state_size = self._optimizer_state_size + var.size
+    logging.info("Total optimizer state: %s", f"{self._optimizer_state_size:,}")
+
   #  _             _
   # | |_ _ __ __ _(_)_ __
   # | __| "__/ _` | | "_ \

kfac_jax/_src/utils/__init__.py

@@ -48,6 +48,8 @@
 del types
 
 # misc
+deserialize_state_tree = misc.deserialize_state_tree
+serialize_state_tree = misc.serialize_state_tree
 to_tuple_or_repeat = misc.to_tuple_or_repeat
 first_dim_is_size = misc.first_dim_is_size
 fake_element_from_iterator = misc.fake_element_from_iterator

kfac_jax/_src/utils/misc.py

@@ -117,17 +117,16 @@ def copy(self):
     (flattened, structure) = jax.tree_util.tree_flatten(self)
     return jax.tree_util.tree_unflatten(structure, flattened)
 
-  def tree_flatten(self) -> Tuple[Tuple[ArrayTree, ...], None]:
-    return self.field_values, None
+  def tree_flatten(self) -> Tuple[Tuple[ArrayTree, ...], Tuple[str, ...]]:
+    return self.field_values, self.field_names()
 
   @classmethod
   def tree_unflatten(
       cls,
-      aux_data: None,
+      aux_data: Tuple[str, ...],
       children: Tuple[ArrayTree, ...],
   ):
-    del aux_data  # not used
-    return cls(**dict(zip(cls.field_names(), children)))
+    return cls(**dict(zip(aux_data, children)))
 
   def __repr__(self) -> str:
     return (f"{self.__class__.__name__}(" +