hk.add_loss about dm-haiku HOT 2 CLOSED

Hey @cgarciae , thanks for the FR! In general we try to keep Haiku fairly lean and encourage features (e.g. training loops, optimizers etc) to be solved in other libraries (then they can benefit all JAX users not just Haiku users) or built from existing Haiku/JAX features.

Wrt using existing features, you might consider using hk.set_state for this. This is a fairly general mechanism in Haiku for logging values associated with modules. You could use hk.data_structures.filter to extract all losses from the state dict:

def f(x):
  y = hk.nets.ResNet50(1000)(x, True)
  loss_1 = y.sum()
  loss_2 = loss_1 ** 2
  hk.set_state("loss_1", loss_1)
  return loss_2

f = hk.transform_with_state(f)

rng = jax.random.PRNGKey(42)
x = jnp.ones([1, 224, 224, 3])
params, state = f.init(rng, x)

# Apply as usual:
params, state = f.apply(params, state, rng, x)

# Extract losses from state:
is_loss = lambda m, n, v: n.startswith("loss")
losses = hk.data_structures.filter(is_loss, state)
print(losses)  # frozendict({'~': frozendict({'loss_1': DeviceArray(0., dtype=float32)})})

If you want to implement this as a standalone feature (e.g. decoupled from hk.set_state) then I've forked the following from @ibab who has implemented something similar (his version is more robust with thread safety and nesting).

from contextlib import contextmanager

loggers = []

@contextmanager
def context():
  data = {}
  loggers.append(data)
  try:
    yield data
  finally:
    loggers.pop()

def log(name, value):
  # NOTE: log(..) ignored when not logging.
  if loggers:
    data = loggers[-1]
    data[name] = value

def f(x):
  x = x ** 2
  log("a", x)
  x = x ** 2
  log("b", x)
  return x

def g():
  with context() as data:
    y = f(2)
  return y, data

y, data = g()
assert y == 16
assert data == {'a': 4, 'b': 16}

I hope that's useful, please feel free to reopen if this does not solve your usecase.

from dm-haiku.

Thanks for the info @tomhennigan ! I think set_state + filters can achieve what I am looking for :)
I didn't realize set_state could be called without a corresponding get_state.

I still think a first class logger would be nice / less error prone but fortunately its not a blocker.

from dm-haiku.