Numpyro: Integrate with flax/haiku

[fehiepsi]

Currently, we support stax but it lacks some popular layers such as RNN. Some other candidates include flax, haiku. It might not take much effort to use those libraries in NumPyro.

We might start with some utilities/handlers in numpyro.contrib.flax or numpyro.contrib.haiku, add some examples (which illustrates how to set priors for some parameters of neural nets), and gradually enhance the API.

Some references:

• lift handler: converts a param site to a sample site
• module primitive: move nn's parameters to param statement
• VAE example

[TuanNguyen27]

Can i work on this ? Promise i'll finish the PR timely :D

[fehiepsi]

Hi @TuanNguyen27! It is great to hear that you are interested in this PR. Please take your time.

FYI, if you come up with several design choices and don't know which one is good, you can create a design doc (you can take a look a some design docs that we created previously) and invite us to review/comment.

[TuanNguyen27]

@fehiepsi For my first attempt, I will try to migrate VAE example to use flax/haiku. If that works, how do I further integrate this library into numpyro, do i re-implement some existing handlers to work with haiku, for example ?

[fehiepsi]

@TuanNguyen27 I think you can put required utilities/handlers in contrib.flax. Users will need to install flax/haiku to use the utilities/handlers in contrib.flax. You can add an ImportError for it, for example like this.

[TuanNguyen27]

@fehiepsi I'm trying to rewrite numpyro.module to work with haiku's Module class, but somehow my losses come out as Nan during training. I'm wondering if you could take a look to see if there's anything i missed.

https://gist.github.com/TuanNguyen27/57ab8264632ccded683ff819acb99253

Runnable colab link here

I had to supply rng keys every time .apply is called, and stax's apply doesn't require this, so I'm not sure how things translate between the two libraries.

[TuanNguyen27]

I'm running into the same problem when using flax's nn.Module. I'd love to hear any suggestion / debugging tips to make more progress :)

https://colab.research.google.com/gist/TuanNguyen27/409192b97d01e1ea5583f9e4499d306a/flax_prototype.ipynb

[TuanNguyen27]

haiku's Module & flax's Module both require user to first initialize the parameter with rng & data input (data input could be completely unrelated to the real training data as long as it has the correct shape). When I reimplement numpyro.module, I initialized Encoder and Decoder's parameters with (rng_key, jnp.ones(input_shape)). I then wrap the initalized params with numpyro.param(module_key, nn_params) -- this is where I am not sure if the code works the same as stax.

[fehiepsi]

@TuanNguyen27 Awesome! You are so fast... Your code looks very reasonable to me. I will try to debug to see if I can isolate the nan issue for you.

[fehiepsi]

@TuanNguyen27 It seems that you need to add some additional layers to the output of those nn modules:

class Encoder(nn.Module):
    def apply(self, x, hidden_size = 512, latents = 10):
        ...
        return mean_x, jnp.exp(logvar_x)

class Decoder(nn.Module):
    def apply(self, z, hidden_size=512, output_shape=784):
        ...
        return nn.sigmoid(z)

While playing with your example, I really like the way you simplify the problem: we only need to make a new handler module and its definition is quite similar to the current numpyro.module. With that, to integrate with NumPyro, I think you only need to put it in a contrib file like numpyro/contrib/nn.py (probably make the name explicitly like module -> stax_module).

The reason for putting it into contrib is I am seeing there are further cool extensions from your work! It is important that those Flax parameters have corresponding names. This enables us to implement something similar to random_module in Pyro, which turns an NN into a Bayesian NN. I remember some users in Pyro forum have asked for this feature. :)

cc @neerajprad I think you will be interested in @TuanNguyen27's work :)

[TuanNguyen27]

that's great! I am also interested in converting a traditional NN into a Bayesian NN. In the past I have used pyvarinf, which might be similar to how Pyro does the conversion. Let me rerun my notebook with your suggestion and make a PR :)

PS: Do you want both haiku and flax to integrate with numpyro or just one ? I can also try to see if Oryx has a similar module system.

[fehiepsi]

converting a traditional NN into a Bayesian NN

Yeah, in Pyro, we first use sample primitives to get the value of some parameters. Then using param primitive for the remaining parameters. I think we can do the same here but need to work around the nested structure of Flax parameter names. I am looking forward to seeing your PR on that. :)

Do you want both haiku and flax to integrate with numpyro or just one ?

I am seeing your PR which includes two. Let's do it. Thanks!

[TuanNguyen27]

@fehiepsi I thought when you extract params from a pytorch module, the dictionary of params is also nested ? How do you overcome that in pyro ?

[fehiepsi]

In PyTorch, you can call module.named_parameters() to get a list of parameters with flatten-using-dot names (e.g. layer1.linear.weight). Those parameters (together with their flatten names) will be stored in a global ParamsStore. Then we use pyro.lift to convert some parameters to random variables.

I think we can use similar approach here. The primitive random_module can be something like this

def random_module(name, nn, prior):
    apply_partial = module(name, nn)
    params = apply_partial.args[0]
    new_params = deepcopy(params)
    # loops over the nested dict `new_params`, get flatten name
    # if it appears in `prior` then we replace its value by a `sample` statement
    # (with name of `sample` is `name` + `$$$` + the flatten name)
    # at the same time, we replace the corresponding value in `params` to `None`
    # to not carry those values around in the trace
    return partial(apply_partial.func, new_params, *apply_partial.args[1:], **apply_partial.keywords)

I think this way, the implementation will be much simpler than Pyro implementation. If you need me clarify any point, please let me know. Here I assumed module returns a partial apply function, with its first arg is params. It would be nice if you modify #717 a bit to get that pattern. :)

[fehiepsi]

@TuanNguyen27 Here is an example for looping over a nested dict and updating corresponding params:

from copy import deepcopy
params = {'a': {'b': {'c': {'d': 1}, 'e': 2}, 'f': 3}}
prior = {'a.b.c.d': 4, 'a.f': 5}
new_params = deepcopy(params)

def update_params(prefix, params, new_params, prior):
    for name, item in params.items():
        flatten_name = ".".join([prefix, name]) if prefix else name
        if isinstance(item, dict):
            assert flatten_name not in prior
            new_item = new_params[name]
            update_params(flatten_name, item, new_item, prior)
        elif not isinstance(prior, dict):
            params[name] = None
            new_params[name] = prior
        elif flatten_name in prior:
            params[name] = None
            # NB: replace prior[flatten_name] by a sample statement from it
            new_params[name] = prior[flatten_name]

update_params('', params, new_params, prior)
assert params == {'a': {'b': {'c': {'d': None}, 'e': 2}, 'f': None}}
assert new_params == {'a': {'b': {'c': {'d': 4}, 'e': 2}, 'f': 5}}

[TuanNguyen27]

@fehiepsi thanks for the detailed code example !

I'm afraid my current implementation for haiku doesn't allow me to access params from a partial function

return lambda x: nn.apply(nn_params, rng_key, x) is a <function haiku_module.<locals>.<lambda> with no attribute. Printing nn.__dict__ returns {}

It is possible to do this for Flax, since the return result of flax_module is flax.nn.Model(nn, nn_params), which has params attribute.

Wondering if there's a work-around to make haiku behave like flax...

Example notebook: https://colab.research.google.com/gist/TuanNguyen27/085c00fa9af777174060430240648449/params.ipynb

[fehiepsi]

Hmm, I think for haiku, you can simply return partial(nn.apply, nn_params, rng_key) and for flax, you can simply return partial(partial(flax.nn.Model, nn), nn_params) or partial(nn.call, nn_params).

[TuanNguyen27]

interesting, so if we wrap .apply with partial we can access nn_params later on ?

[fehiepsi]

Yes, you can see how it works in this comment.

[fehiepsi]

@TuanNguyen27 I think we have all ingredients for random_module. Do you want to pair-code this feature together? In case you agree, could you setup some time this weekend or next week (I believe 30m or 1h is enough)? I would like to hear more about your experience with pyvarinf and your ideas on a good example for random_module. :)

[TuanNguyen27]

Yes, i've been wanting to work on this but I forgot, I will ping you directly to set up a time :)

[quattro]

I realized this is closed, and already implemented, but I'm curious if it makes sense to have 'sister' instantiating functions for haiku_modules that take *args, **kwargs rather than input shape. This would enable more complex haiku modules with multiple parameters. For example,

def haiku_general_module(name, nn_module, *args, **kwargs):
    module_key = name + '$params'
    nn_params = numpyro.param(module_key)
    if nn_params is None:

        # feed in dummy data to init params
        rng_key = numpyro.prng_key()
        nn_params = nn_module.init(rng_key, *args, **kwargs)

        # haiku init returns an immutable dict
        nn_params = hk.data_structures.to_mutable_dict(nn_params)

        # we cast it to a mutable one to be able to set priors for parameters
        # make sure that nn_params keep the same order after unflatten
        params_flat, tree_def = tree_flatten(nn_params)
        nn_params = tree_unflatten(tree_def, params_flat)
        numpyro.param(module_key, nn_params)

    return ft.partial(nn_module.apply, nn_params, None)

This should cleanly translate to enable situations like the following:

dual_enc_f = hk.transform(lambda x_0, x_1: DualEncoder(hidden_dim, latent_dim, activate_f)(x_0, x_1))
encoder_m = haiku_general_module("encoder", dual_enc_f, x_0 = x_0_data, x_1 = x_1_data)
z_loc, z_std = encoder_m(x_0_data, x_1_data)

[fehiepsi]

Hi @quattro, please feel free to make a pull request to make haiku_module (and flax_module too) more flexible. I think you don't need to implement a new function. How about

def haiku_module(name, nn_module, *, input_shape=None, **kwargs):
    ...
    if input_shape is not None:
        args = (jnp.ones(input_shape),) if input_shape is not None else ()
    nn_params = nn_module.init(rng_key, *args, **kwargs)

?