Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ResBlock combinator in stax #1569

Closed
wants to merge 3 commits into from
Closed

ResBlock combinator in stax #1569

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6bdc3b2
ResBlock combinator in stax
aeftimia Oct 2, 2019
05b891e
change default tail to Identity layer
aeftimia Oct 28, 2019
9094716
python2 compatibility
aeftimia Oct 28, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions examples/resnet50.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,9 +30,9 @@
from jax import jit, grad, random
from jax.experimental import optimizers
from jax.experimental import stax
from jax.experimental.stax import (AvgPool, BatchNorm, Conv, Dense, FanInSum,
FanOut, Flatten, GeneralConv, Identity,
MaxPool, Relu, LogSoftmax)
from jax.experimental.stax import (AvgPool, BatchNorm, Conv, Dense,
Flatten, GeneralConv, Identity,
MaxPool, Relu, ResBlock, LogSoftmax)


# ResNet blocks compose other layers
Expand All @@ -45,7 +45,7 @@ def ConvBlock(kernel_size, filters, strides=(2, 2)):
Conv(filters2, (ks, ks), padding='SAME'), BatchNorm(), Relu,
Conv(filters3, (1, 1)), BatchNorm())
Shortcut = stax.serial(Conv(filters3, (1, 1), strides), BatchNorm())
return stax.serial(FanOut(2), stax.parallel(Main, Shortcut), FanInSum, Relu)
return ResBlock(Main, Shortcut)


def IdentityBlock(kernel_size, filters):
Expand All @@ -58,7 +58,7 @@ def make_main(input_shape):
Conv(filters2, (ks, ks), padding='SAME'), BatchNorm(), Relu,
Conv(input_shape[3], (1, 1)), BatchNorm())
Main = stax.shape_dependent(make_main)
return stax.serial(FanOut(2), stax.parallel(Main, Identity), FanInSum, Relu)
return ResBlock(Main, Identity)


# ResNet architectures compose layers and ResNet blocks
Expand Down
16 changes: 16 additions & 0 deletions jax/experimental/stax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -260,6 +260,22 @@ def apply_fun(params, inputs, **kwargs):
return init_fun, apply_fun


def ResBlock(*layers, fan_in=FanInSum, tail=Relu):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can’t use this syntax yet because JAX still supports Python 2 (for a little while longer)

Copy link
Contributor Author

@aeftimia aeftimia Oct 28, 2019 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks reasonable to me, though my preference would be to avoid the final “tail” layer defaulting to ReLU and keep all activations explicit. This would be more consistent with the rest of stax, and in my experience unexpected activations are a common source of bugs.

Oh interesting--in hindsight I could see that potentially being an issue but I don't think I would have thought of it before you mentioned it. I modified the default tail to be Identity.

You can’t use this syntax yet because JAX still supports Python 2 (for a little while longer)

I submitted a patch that works for Python 2 and 3 with a TODO to replace it with the more readable function head when Python 2 support is dropped.

Please feel free to let me know if you would like any further modifications.

"""Split input, feed it through one or more layers in parallel,
recombine them with a fan-in, apply a trailing layer (i.e. an activation)

Args:
*layers: a sequence of layers, each an (init_fun, apply_fun) pair.
fan_in, optional: a fan-in to recombine the outputs of each layer
tail, optional: a final layer to apply after recombination

Returns:
A new layer, meaning an (init_fun, apply_fun) pair, representing the
parallel composition of the given sequence of layers fed into fan_in and then tail.
"""
return serial(FanOut(len(layers)), parallel(*layers), fan_in, tail)


# Composing layers via combinators


Expand Down