fix physionet

latent_ode.py

@@ -16,7 +16,7 @@
 from jax.experimental.jet import jet
 from jax.flatten_util import ravel_pytree
 
-import lib
+from lib.optimizers import exponential_decay
 from lib.ode import odeint
 from physionet_data import init_physionet_data
 
@@ -57,7 +57,6 @@
 rng = jax.random.PRNGKey(seed)
 dirname = parse_args.dirname
 count_nfe = not parse_args.no_count_nfe
-num_blocks = parse_args.num_blocks
 ode_kwargs = {
  "atol": parse_args.atol,
  "rtol": parse_args.rtol
@@ -276,31 +275,31 @@ def init_model(gen_ode_kwargs,
  "b_init": jnp.zeros
  }
 
- gru_rnn = hk.transform(wrap_module(LatentGRU,
- latent_dim=rec_dim,
- n_units=gru_units,
- **init_kwargs))
+ gru_rnn = hk.without_apply_rng(hk.transform(wrap_module(LatentGRU,
+  latent_dim=rec_dim,
+  n_units=gru_units,
+  **init_kwargs)))
  gru_rnn_params = gru_rnn.init(rng, *initialization_data_["gru_rnn"])
 
  # note: the ODE-RNN version uses double
- rec_to_gen = hk.transform(wrap_module(lambda: hk.Sequential([
+ rec_to_gen = hk.without_apply_rng(hk.transform(wrap_module(lambda: hk.Sequential([
  lambda x, y: jnp.concatenate((x, y), axis=-1),
  hk.Linear(50, **init_kwargs),
  jnp.tanh,
  hk.Linear(2 * gen_dim, **init_kwargs)
- ])))
+ ]))))
  rec_to_gen_params = rec_to_gen.init(rng, *initialization_data_["rec_to_gen"])
 
- gen_dynamics = hk.transform(wrap_module(GenDynamics,
- latent_dim=gen_dim,
- units=dynamics_units,
- layers=gen_layers))
+ gen_dynamics = hk.without_apply_rng(hk.transform(wrap_module(GenDynamics,
+  latent_dim=gen_dim,
+  units=dynamics_units,
+  layers=gen_layers)))
  gen_dynamics_params = gen_dynamics.init(rng, *initialization_data_["gen_dynamics"])
  gen_dynamics_wrap = lambda x, t, params: gen_dynamics.apply(params, x, t)
 
- gen_to_data = hk.transform(wrap_module(hk.Linear,
- output_size=data_dim,
- **init_kwargs))
+ gen_to_data = hk.without_apply_rng(hk.transform(wrap_module(hk.Linear,
+  output_size=data_dim,
+  **init_kwargs)))
  gen_to_data_params = gen_to_data.init(rng, initialization_data_["gen_to_data"])
 
  init_params = {
@@ -310,7 +309,7 @@ def init_model(gen_ode_kwargs,
  "gen_to_data": gen_to_data_params
  }
 
- def forward(count_nfe_, reg, _method, params, data, data_timesteps, timesteps, mask, num_samples=3):
+ def forward(count_nfe_, params, data, data_timesteps, timesteps, mask, num_samples=3):
  """
  Forward pass of the model.
  y are the latent variables of the recognition model
@@ -343,7 +342,7 @@ def integrate_sample(z0_):
  dynamics = gen_dynamics_wrap
  init_fn = lambda x: x
  else:
- dynamics = augment_dynamics(gen_dynamics_wrap, reg)
+ dynamics = augment_dynamics(gen_dynamics_wrap)
  init_fn = aug_init
  return jax.vmap(lambda z_, t_: odeint(dynamics, init_fn(z_), t_,
  params["gen_dynamics"], **gen_ode_kwargs),
@@ -391,7 +390,7 @@ def scan_fun(prev_state, xi):
  model = {
  "forward": partial(forward, False),
  "params": init_params,
- "nfe": lambda *args: partial(forward, count_nfe, reg)(*args)[-1]
+ "nfe": lambda *args: partial(forward, count_nfe)(*args)[-1]
  }
 
  return model
@@ -474,10 +473,10 @@ def run():
  forward = lambda *args: model["forward"](*args)[1:]
  grad_fn = jax.grad(lambda *args: loss_fn(forward, *args))
 
- lr_schedule = lib.optimizers.exponential_decay(step_size=parse_args.lr,
-  decay_steps=1,
-  decay_rate=0.999,
-  lowest=parse_args.lr / 10)
+ lr_schedule = exponential_decay(step_size=parse_args.lr,
+ decay_steps=1,
+ decay_rate=0.999,
+ lowest=parse_args.lr / 10)
  opt_init, opt_update, get_params = optimizers.adamax(step_size=lr_schedule)
  opt_state = opt_init(model["params"])
 
@@ -574,7 +573,7 @@ def evaluate_loss(opt_state, ds_test, kl_coef):
 
  print(print_str)
 
- outfile = open("%s/reg_%s_lam_%.12e_num_blocks_%d_info.txt" % (dirname, reg, lam, num_blocks), "a")
+ outfile = open("%s/reg_%s_lam_%.12e_info.txt" % (dirname, reg, lam), "a")
  outfile.write(print_str + "\n")
  outfile.close()
 
@@ -594,14 +593,14 @@ def evaluate_loss(opt_state, ds_test, kl_coef):
  pickle.dump(fargs, outfile)
  outfile.close()
 
- outfile = open("%s/reg_%s_lam_%.12e_num_blocks_%d_iter.txt" % (dirname, reg, lam, num_blocks), "a")
+ outfile = open("%s/reg_%s_lam_%.12e_iter.txt" % (dirname, reg, lam), "a")
  outfile.write("Iter: {:04d}\n".format(itr))
  outfile.close()
  meta = {
  "info": info,
  "args": parse_args
  }
- outfile = open("%s/reg_%s_lam_%.12e_num_blocks_%d_meta.pickle" % (dirname, reg, lam, num_blocks), "wb")
+ outfile = open("%s/reg_%s_lam_%.12e_meta.pickle" % (dirname, reg, lam), "wb")
  pickle.dump(meta, outfile)
  outfile.close()