Added max and avg pool, alongside tests and docs #77
Conversation
There was a problem hiding this comment.
Yes, I'd prefer to avoid the additional dependency on einops. Good thinking to have the conv/pooling layers use the same dimension ordering. (Honestly the lax.conv_{...} options are a mess of conflicting defaults.)
Can you add some tests for the backward pass? As per #66 it seems like that's a potential issue we may need to try and resolve.
|
Also: I think the new pooling documentation needs to be added to |
equinox/nn/pool.py
Outdated
| class Pool(Module): | ||
| """ General N-dimensional downsampling over a sliding window. """ | ||
|
|
||
| operation: Callable[[Array], Array] = static_field() |
There was a problem hiding this comment.
This shouldn't be static. (In principle the operation may be learnt.) Also I think the annotation should be Callable[[Array, Array], Array].
equinox/nn/pool.py
Outdated
|
|
||
| def __init__( | ||
| self, | ||
| operation: Callable[[Array], Array], |
equinox/nn/pool.py
Outdated
| self.operation = operation | ||
| self.init = 0.0 | ||
| if operation is lax.max: | ||
| self.init = -jnp.inf |
There was a problem hiding this comment.
Make init an input, don't try to guess based on special values. Essentially Pool should make sense as a standalone operation.
equinox/nn/pool.py
Outdated
| if isinstance(padding, int): | ||
| self.padding = tuple((padding, padding) for _ in range(num_spatial_dims)) | ||
| elif isinstance(padding, Sequence) and \ | ||
| all(isinstance(element, tuple) for element in padding): |
There was a problem hiding this comment.
isinstance(element, tuple)
to
isinstance(element, Sequence)
|
Made changes, think I have correctly typeset to the flake8 precommit, but I don't have approval to run workflows. |
|
Also with regards to the non-differentiability, my understanding is that there was a small mistake in #66 with the init for the maximum pooling. From the tensorflow docs `All initial values have to form an identity under computation' (https://www.tensorflow.org/xla/operation_semantics#reducewindow). Hence when using lax.add as your computation you need to init the pooling at 0.0, as lax.add of 0 and your first element will be your first element. Similarly when using lax.max, you need to init at -jnp.inf, as lax.max of -jnp.inf and your first element will always be your first element. This has been noted to be flagged as the error "Differentiation rule for 'reduce_window' not implemented" in Jax, google/jax#7718. |
equinox/nn/pool.py
Outdated
| kernel_size: Union[int, Sequence[int]], | ||
| stride: Union[int, Sequence[int]] = 1, | ||
| padding: Union[int, Sequence[int], Sequence[tuple[int, int]]] = 0, | ||
| init: int = 0, |
There was a problem hiding this comment.
I think the choice of init is intrinsically tied to the operation. So I wouldn't give this any default value at all, and instead demand that it be user-specified. (For consistency with lax.reduce_window I'd make the argument order init then operation, by the way.)
equinox/nn/pool.py
Outdated
| self.operation = operation | ||
| self.init = init | ||
| if operation is lax.max: | ||
| self.init = -jnp.inf |
There was a problem hiding this comment.
This line still needs to be removed.
equinox/nn/pool.py
Outdated
| num_spatial_dims: int, | ||
| kernel_size: Union[int, Sequence[int]], | ||
| stride: Union[int, Sequence[int]] = 1, | ||
| padding: Union[int, Sequence[int], Sequence[tuple[int, int]]] = 0, |
There was a problem hiding this comment.
This needs to be typing.Tuple, not tuple. The latter only works for Python 3.9+.
equinox/nn/pool.py
Outdated
| - `key`: Ignored; provided for compatibility with the rest of the Equinox API. | ||
| (Keyword only argument.) | ||
| **Returns:** | ||
| A JAX array of shape `(in_channels, new_dim_1, ..., new_dim_N)`. |
There was a problem hiding this comment.
Nit: rename both instances of in_channels to just channels, since there's no in/out distinction here.
tests/test_nn.py
Outdated
| @@ -6,6 +6,8 @@ | |||
| import jax.random as jrandom | |||
| import pytest | |||
|
|
|||
| import sys | |||
| sys.path.append('/Users/benwalker/PycharmProjects/equinox') | |||
There was a problem hiding this comment.
Hahahah oops, didn't mean to leave that one in
|
Various nits above but I think we're nearly there. Nice work on the nondifferentiability; I'm glad to hear that won't be an issue. |
|
Made those changes, also for the differentiability do you think it is worth noting somewhere (maybe in the arguments docstring?) that if you use Pool by itself you need to make sure that operation(init, x) = x for all finite x? |
equinox/nn/pool.py
Outdated
|
|
||
| def __init__( | ||
| self, | ||
| init: int, |
There was a problem hiding this comment.
Nit: not necessarily an integer! e.g. jnp.NINF is a float.
Admittedly I'm not sure what the correct annotation is here. Union[int, float] maybe? Probably Arrays are valid too?
There was a problem hiding this comment.
I have gone for Union[int, float, Array], haven't seen anyone use an Array but if the elements you were comparing are arrays then this would be necessary.
|
Fixed the remaining typesetting mistake |
|
@Benjamin-Walker Checks have failed again! By the way, to tidy up this branch before merging, can you squash all the commits together? I think the appropriate commands are something like (untested): Noting that the last command will force-push and thus override this branch. (If you're feeling uncertain about that then you can use a different branch name, and open a new PR, instead.) |
|
I think it was actually 3.10 where they introduced using | instead of Union. Think I have squashed all the commits together, let me know if it hasn't worked. |
* Improved filtered transformations (#71) * Updated filter_{jit,grad,value_and_grad} * Documented that BatchNorm supports multiple axis_names * Doc fix * added tree_inference * Added filter_vmap and filter_pmap * Made filter_{vmap,pmp} non-experimental. Fixed doc issues. * Added 'inference' call-time argument to 'MultiheadAttention', which was missed when switching over 'Dropout' to use 'inference'. * Added test for tree_inference * Doc improvements * Doc tweaks * Updated tests * Two minor bugfixes for filter_{jit,grad,vmap,pmap,value_and_grad}. 1. Applying filter_{jit,vmap,pmap} to a function with *args and **kwargs should now work. Previously inspect.Signature.apply_defaults was not filling these in. 2. The output of all filtered transformations has become a Module. This means that e.g. calling filter_jit(filter_grad(f)) multiple times will not induce any recompilation, as all filter_grad(f)s are PyTrees of the same structure as each other. * Crash and test fixes * tidy up * Fixes and test fixes * Finished adding tests. * added is_leaf argument to filter and partition (#72) * Improved tree_at a lot: (#73) - Can now substitute arbitrary nodes, not just leaves - These substituted nodes can now have different structures to each other. - Will raise an error if `where` depends on the values of the leaves of the PyTree. In addition, `tree_equal` should now work under JIT. * Lots of doc fixes and tweaks (#74) * Added tree_{de,}serialise_leaves (#80) * Added max/average pooling. (#77) * Tidy up pooling implementation (#81) * Version bump * Fixed tree_at breaking when replacing Nones. (#83) Co-authored-by: Ben Walker <38265558+Benjamin-Walker@users.noreply.github.com>
I currently use einops.rearrange so that the pooling layers and conv layers take the same tensor format as input, BCHW. For some reason lax.conv_general_dilated and lax.reduce_window are BCHW and BHWC respectively. Can do this not using einops if you would prefer to not have the extra dependency.