How do I load a pre-trained model?

[eleninisioti]

There is a notebook that explains how to save and load models (https://github.com/google/brax/blob/main/notebooks/training.ipynb) but there testing happens right after training, calling function make_inference_fn(params), which requires first running make_inference_fn, params, _ = train_fn(environment=env, progress_fn=progress), the whole training process.

My question is how I can test without running the training process, simply by loading the params. How can I have inference_fn withour training?

[btaba]

Hi @eleninisioti ! This is a great feature request that we haven't gotten around to implement – A model saver that saves the parameter config alongside the weights.

Here is an example of what the loader would do (where load_config_dict is to-be-implemented, and the whole thing needs to be wrapped in a nice function):

def make_inference_fn(
    observation_size: int,
    action_size: int,
    normalize_observations: bool = True,
    network_factory_kwargs: Optional[Dict[str, Any]] = None,
):
  normalize = lambda x, y: x
  if normalize_observations:
    normalize = running_statistics.normalize
  ppo_network = brax_networks.make_ppo_networks(
      observation_size,
      action_size,
      preprocess_observations_fn=normalize,
      **(network_factory_kwargs or {}),
  )
  make_policy = brax_networks.make_inference_fn(ppo_network)
  return make_policy

config_dict = load_config_dict(checkpoint_path)
make_policy = make_inference_fn(
    config_dict['observation_size'],
    config_dict['action_size'],
    config_dict['normalize_observations'],
    network_factory_kwargs=config_dict['network_factory_kwargs'], # {"policy_hidden_layer_sizes": (128,) * 4}
)
params = model.load_params(checkpoint_file)
jit_inference_fn = jax.jit(make_policy(params, deterministic=True))

[m-serra]

Hi!
Any updates on this?
Without the ability to save and load inference_fn, or at least re-creating it, how can someone do evaluation on a model checkpoint?
Maybe this should be an Issue...
Thank you in advance!

[btaba]

Created an issue #438

You can use the above example for now where you hard-code values for observation_size etc.

[arcman7]

I'm a bit of a newb here... I haven't quite figured out yet how to resume training from some previously trained model params that I would load at the start of the experiment; how can I pass in pre-trained model params before running the train_fn function?

[AlexS28]

would be great to know this as well, so we can load a pre-trained model and continue training from that model.

[arcman7]

@AlexS28 Hey, we got something working, a bit hacky but it works - I'll ask @jihan1218 to post a link here to his solution

[AlexS28]

that would be great thanks!

[jihan1218]

hey guys I made a solution for this. https://github.com/jihan1218/brax
this will allow you to feed your previously trained parameters and continue your experiment!

[AlexS28]

thanks!! looks great.

[btaba]

Please see #438

This makes checkpointing available with minimal changes to the public API.