Implement Transfer Learning API

[n2cholas]

Closes #1. Implements a Sequential combinator as well as a slice to easily extract portions of the models.

TODO:

• Write more thorough tests
• Write documentation

[n2cholas]

@cgarciae any thoughts on this initial design?

[cgarciae]

Hey @n2cholas ! We had a discussion in this PR as ways to approach this: poets-ai/elegy#169

My thoughts:

• Slicing is simple and probably the way to go with Flax.
• elegy.Module enables you to do Transfer Learning in more flexible and easier fashion but I don't want push it too hard because the ecosystem seems to be unifying around Flax. I might write a blog with the approach we use.

[n2cholas]

Elegy's approach looks very flexible, I'll definitely give it a try in the future. A blog post would be very nice!

Since this project's scope is ResNet-style architectures, I'll iterate on this slice API a bit and stick to it.

I'm curious to see how Flax will address this problem. nn.compact gives a great boost for ease of implementation and readability, but somewhat sacrifices re-usability of the code (since, for now, you have to treat the Module like black box).

I actually quite liked jax.experimental.stax's combinator ideology--it made dataflow explicit and would enable simple arbitrary model surgery.

[n2cholas]

@cgarciae we'll provide a Sequential module and a slice_variables method. It's trivial to slice a sequential model yourself (sliced_model = Sequential(model.layers[start:end]), and slice_variables will give you the corresponding variables dict (sliced_variables = slice_variables(variables, start, end)).

Once Flax has its own Sequential module (PR), I'll switch to that.

Thanks again for opening this up!