# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved. # # 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 # # http://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. """K-FAC for accumulating statistics.""" from typing import Any, Callable, Generic import jax import jax.numpy as jnp from kfac_jax._src.utils import misc from kfac_jax._src.utils import parallel from kfac_jax._src.utils import types Array = types.Array Numeric = types.Numeric Shape = types.Shape DType = types.DType ArrayTree = types.ArrayTree TArrayTree = types.TArrayTree AddFunction = Callable[[TArrayTree, TArrayTree, Numeric, Numeric], TArrayTree] def default_add_function( obj1: TArrayTree, obj2: TArrayTree, coeff1: Numeric, coeff2: Numeric ) -> TArrayTree: return jax.tree_util.tree_map( lambda x, y: coeff1 * x + coeff2 * y, obj1, obj2) @misc.register_state_class class WeightedMovingAverage(Generic[TArrayTree], misc.State): """A wrapped class for an arbitrary weighted moving average.""" weight: Numeric value: TArrayTree | None @property def ndim(self) -> int: assert self.value is not None return self.value.ndim # pytype: disable=attribute-error @property def shape(self) -> Shape: assert self.value is not None return self.value.shape # pytype: disable=attribute-error @property def dtype(self) -> DType: assert self.value is not None return self.value.dtype # pytype: disable=attribute-error def update( self, value: TArrayTree, old_weight_multiplier: Numeric, new_weight: Numeric, add_function: AddFunction = default_add_function, ): """Updates the underlying array and weight accordingly.""" assert self.value is not None # A negative value of new_weight means we should only update the value # (with -new_weight) and not the total running weight. This roughly # corresponds to summation instead of averaging, and is useful in a few # contexts. self.weight = old_weight_multiplier * self.weight + jax.nn.relu(new_weight) eta_for_old = jax.nn.relu(new_weight) / self.weight eta_for_new = jnp.abs(new_weight) / self.weight self.value = add_function(self.value, value, 1.0 - eta_for_old, eta_for_new) def sync(self, pmap_axis_name: str | None): """Syncs the underlying array across devices.""" if self.value is None: raise ValueError("`_value` has not been set yet.") self.value = parallel.pmean_if_pmap(self.value, pmap_axis_name) def clear(self, value_to_none: bool = False): """Resets the weighted average.""" self.weight = jnp.zeros_like(self.weight) self.value = None if value_to_none else jnp.zeros_like(self.value) def value_and_clear(self) -> TArrayTree: """Retrieves the value of the weighted average and clears it.""" value = self.value self.clear() assert value is not None return value @classmethod def zeros_array( cls, shape: Shape, dtype: DType | None = None, ) -> "WeightedMovingAverage[Array]": """Initializes a `WeightedMovingAverage` with a single array of zeros.""" return cls( # pytype: disable=wrong-keyword-args weight=jnp.zeros([], dtype=dtype), value=jnp.zeros(shape, dtype=dtype), ) @classmethod def zeros_like(cls, value: TArrayTree) -> "WeightedMovingAverage[TArrayTree]": """Initializes a `WeightedMovingAverage` with zeros structure like `value`.""" return cls( # pytype: disable=wrong-keyword-args weight=jnp.array( 0.0, dtype=types.get_float_dtype_and_check_consistency(value) ), value=jax.tree_util.tree_map(jnp.zeros_like, value), ) class MultiChunkAccumulator(Generic[TArrayTree]): """Statistics accumulation, abstracted over multiple chunks.""" def __init__( self, init_obj_value: TArrayTree | None, weight: Numeric, multi_device: bool, ): """Initializes an accumulator instance with the provided object and counter. Args: init_obj_value: The initial value of the accumulator. weight: The initial weight, which specifies how many samples are assumed to have been already counted in the initial value of the accumulator. multi_device: Whether the objects that are accumulated are outputs of a multi-device computation (e.g. `jax.pmap`). """ self._accumulator = init_obj_value self._weight = weight self._multi_device = multi_device @property def accumulator(self) -> TArrayTree | None: """The current value of the underlying not-normalized accumulator.""" return self._accumulator @property def weight(self) -> Numeric | None: """The current normalization weight of the underlying accumulator.""" return self._weight @property def multi_device(self) -> bool: """Whether the accumulator is the output of a multi-device computation.""" return self._multi_device @property def value(self) -> TArrayTree | None: """The current normalized value of the accumulator.""" if types.tree_is_empty(self.accumulator): return self.accumulator if self._multi_device: return parallel.pmap_sync_and_divide_value(self.accumulator, self.weight) else: return parallel.jit_sync_and_divide_value(self.accumulator, self.weight) def clear(self) -> None: """Sets the underlying accumulator and weight to `None`.""" self._accumulator = None self._weight = None def value_and_clear(self) -> TArrayTree | None: """Retrieves the normalized value of the accumulator and clears it.""" value = self.value self.clear() return value def add(self, value_obj: TArrayTree, weight: Numeric = 1): """Adds an element to the moving average and the max. The exact update equation for the statistics are: raw_value_t = raw_value_{t-1} + value_obj * weight weight_t = weight_{t-1} + weight Args: value_obj: The value of the object, which scaled by `weight` will be added to the accumulator. weight: The relative weight of the `value_obj`. """ value_obj = jax.tree_util.tree_map(lambda x: x * weight, value_obj) if self._accumulator is None: self._accumulator = value_obj if isinstance(weight, types.SCALAR_TYPES): self._weight = jnp.full_like(self._weight, weight) elif not isinstance(weight, jax.Array): raise ValueError("`weight` should be an instance of float, int or " "jax.Array.") elif self._weight.shape != weight.shape: # pytype: disable=attribute-error # numpy-scalars raise ValueError("If `weight` is an `jnp.ndarray` then should have the " "same shape as the weight of the accumulator.") else: self._weight = weight return if not types.tree_is_empty(self._accumulator): if types.tree_is_empty(value_obj): raise ValueError("The provided `value_obj` has an empty PyTree " "structure, but the accumulator has been initialized " "with a non-empty PyTree object.") self._accumulator = jax.tree_util.tree_map( jnp.add, self._accumulator, value_obj) elif not types.tree_is_empty(value_obj): raise ValueError("The provided `value_obj` has a non-empty PyTree " "structure, but the accumulator has been initialized " "with an empty PyTree object.") self._weight = self._weight + weight @classmethod def zeros_like( cls, obj: TArrayTree, multi_device: bool ) -> "MultiChunkAccumulator[TArrayTree]": """Creates a zero initialized accumulator as `obj`.""" if multi_device: value = (parallel.pmap_zeros_like(obj) if not types.tree_is_empty(obj) else obj) weight = parallel.replicate_all_local_devices( jnp.zeros([], dtype=jnp.int32)) else: value = (parallel.jit_zeros_like(obj) if not types.tree_is_empty(obj) else obj) weight = jnp.zeros([], dtype=jnp.int32) return cls(value, weight, multi_device) @classmethod def empty(cls, multi_device: bool) -> "MultiChunkAccumulator[Any]": """Creates an empty accumulator.""" weight = jnp.zeros([], dtype=jnp.int32) if multi_device: weight = parallel.replicate_all_local_devices(weight) return cls(None, weight, multi_device) def __repr__(self): return (f"{self.__class__.__name__}({self._accumulator!r}, " f"{self._weight!r}, {self._multi_device})") def copy(self): """Returns a copy of the PyTree structure (but not the JAX arrays).""" (flattened, structure) = jax.tree_util.tree_flatten(self) return jax.tree_util.tree_unflatten(structure, flattened) jax.tree_util.register_pytree_node( MultiChunkAccumulator, lambda x: ((x.accumulator, x.weight), (x.multi_device,)), lambda fixed, arrays: MultiChunkAccumulator(*arrays, *fixed) )