Adding "count_include_pad" argument to flax.linen.pooling.avg_pool by dslisleedh · Pull Request #2451 · google/flax

dslisleedh · 2022-09-09T08:30:59Z

What does this PR do?

Now version's flax.linen.pooling.avg_pool average window_sum result include padded tokens. I add argument whether to include padded tokens or not

Checklist

This change is discussed in a Github issue/
issues
The documentation and docstrings adhere to the
documentation guidelines.
This change includes necessary high-coverage tests.
(No quality testing = no merge!)

codecov-commenter · 2022-09-09T12:55:25Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 78.83%. Comparing base (8687673) to head (d4fa16d).
Report is 1820 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2451      +/-   ##
==========================================
- Coverage   79.66%   78.83%   -0.84%     
==========================================
  Files          49       49              
  Lines        4982     5070      +88     
==========================================
+ Hits         3969     3997      +28     
- Misses       1013     1073      +60

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

cgarciae · 2022-09-13T13:59:35Z

Hey @dslisleedh, thanks for creating this PR!
I think this looks good, however, before merging we should probably add a test using this flag.

dslisleedh · 2022-09-13T22:34:46Z

To @cgarciae,

Sorry, I tested my self but didn't share result in PR. Here's result.

for i, shape in enumerate([(1, 5, 3), (1, 5, 5, 3), (1, 5, 5, 5, 3)]):
    inputs = jnp.ones(shape=shape)
    for kernel_size in [(1,), (3,), (5,)]:
        k_s = kernel_size * (i + 1)
        for strides in [(1,), (2,), (3,)]:
            s = strides * (i + 1)
            for padding in ['VALID','SAME']:
                res = avg_pool(inputs, k_s, strides=s, padding=padding, count_include_pad=False)

                print(jnp.min(res), jnp.max(res))

1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0
1.0 1.0

This is first time for me to create PR, so tell me anything if I missed something. Thank you.

cgarciae · 2022-09-14T15:12:48Z

Just noticed there aren't any tests for pooling.py, I'll create an issue, we can add a test for this in the future.

dslisleedh · 2022-09-14T16:31:11Z

I found PoolTest class from ./tests/linen/test_linen.py

class PoolTest(absltest.TestCase):

  def test_pool_custom_reduce(self):
    x = jnp.full((1, 3, 3, 1), 2.)
    mul_reduce = lambda x, y: x * y
    y = nn.pooling.pool(x, 1., mul_reduce, (2, 2), (1, 1), 'VALID')
    np.testing.assert_allclose(y, np.full((1, 2, 2, 1), 2. ** 4))

  def test_avg_pool(self):
    x = jnp.full((1, 3, 3, 1), 2.)
    pool = lambda x: nn.avg_pool(x, (2, 2))
    y = pool(x)
    np.testing.assert_allclose(y, np.full((1, 2, 2, 1), 2.))
    y_grad = jax.grad(lambda x: pool(x).sum())(x)
    expected_grad = jnp.array([
        [0.25, 0.5, 0.25],
        [0.5, 1., 0.5],
        [0.25, 0.5, 0.25],
    ]).reshape((1, 3, 3, 1))
    np.testing.assert_allclose(y_grad, expected_grad)

  def test_avg_pool_no_batch(self):
    x = jnp.full((3, 3, 1), 2.)
    pool = lambda x: nn.avg_pool(x, (2, 2))
    y = pool(x)
    np.testing.assert_allclose(y, np.full((2, 2, 1), 2.))
    y_grad = jax.grad(lambda x: pool(x).sum())(x)
    expected_grad = jnp.array([
        [0.25, 0.5, 0.25],
        [0.5, 1., 0.5],
        [0.25, 0.5, 0.25],
    ]).reshape((3, 3, 1))
    np.testing.assert_allclose(y_grad, expected_grad)

  def test_max_pool(self):
    ...

When I ran this code with the edits from this PR there were no problems.

(flax_test) idongheon@idongheon-ui-MacBookAir flax % python ./tests/linen/linen_test.py 
Running tests under Python 3.10.5: /Users/idongheon/miniforge3/envs/flax_test/bin/python
[ RUN      ] IdsTest.test_hashable
[       OK ] IdsTest.test_hashable
[ RUN      ] NormalizationTest.test_batch_norm
I0915 01:20:42.378859 4314512704 xla_bridge.py:169] Remote TPU is not linked into jax; skipping remote TPU.
I0915 01:20:42.379297 4314512704 xla_bridge.py:345] Unable to initialize backend 'tpu_driver': Could not initialize backend 'tpu_driver'
I0915 01:20:42.379443 4314512704 xla_bridge.py:345] Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I0915 01:20:42.379549 4314512704 xla_bridge.py:345] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I0915 01:20:42.380259 4314512704 xla_bridge.py:345] Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'
[       OK ] NormalizationTest.test_batch_norm
[ RUN      ] NormalizationTest.test_batch_norm_complex
[       OK ] NormalizationTest.test_batch_norm_complex
[ RUN      ] NormalizationTest.test_batch_norm_multi_init
[       OK ] NormalizationTest.test_batch_norm_multi_init
[ RUN      ] NormalizationTest.test_group_norm
[       OK ] NormalizationTest.test_group_norm
[ RUN      ] NormalizationTest.test_group_norm_raises
[       OK ] NormalizationTest.test_group_norm_raises
[ RUN      ] NormalizationTest.test_layer_norm0 (reduction_axes=-1)
[       OK ] NormalizationTest.test_layer_norm0 (reduction_axes=-1)
[ RUN      ] NormalizationTest.test_layer_norm1 (reduction_axes=1)
[       OK ] NormalizationTest.test_layer_norm1 (reduction_axes=1)
[ RUN      ] NormalizationTest.test_layer_norm2 (reduction_axes=(1, 2))
[       OK ] NormalizationTest.test_layer_norm2 (reduction_axes=(1, 2))
[ RUN      ] PoolTest.test_avg_pool
[       OK ] PoolTest.test_avg_pool
[ RUN      ] PoolTest.test_avg_pool_no_batch
[       OK ] PoolTest.test_avg_pool_no_batch
[ RUN      ] PoolTest.test_max_pool
[       OK ] PoolTest.test_max_pool
[ RUN      ] PoolTest.test_pool_custom_reduce
[       OK ] PoolTest.test_pool_custom_reduce
[ RUN      ] RecurrentTest.test_complex_input_gru
[       OK ] RecurrentTest.test_complex_input_gru
[ RUN      ] RecurrentTest.test_convlstm
[       OK ] RecurrentTest.test_convlstm
[ RUN      ] RecurrentTest.test_gru
[       OK ] RecurrentTest.test_gru
[ RUN      ] RecurrentTest.test_lstm
[       OK ] RecurrentTest.test_lstm
[ RUN      ] RecurrentTest.test_optimized_lstm_cell_matches_regular
/Users/idongheon/miniforge3/envs/flax_test/lib/python3.10/site-packages/jax/test_util.py:44: FutureWarning: jax.test_util.check_eq is deprecated and will soon be removed.
  warnings.warn(f"jax.test_util.{attr} is deprecated and will soon be removed.", FutureWarning)
[       OK ] RecurrentTest.test_optimized_lstm_cell_matches_regular
[ RUN      ] StochasticTest.test_dropout
[       OK ] StochasticTest.test_dropout
[ RUN      ] StochasticTest.test_dropout_rate_limits
[       OK ] StochasticTest.test_dropout_rate_limits
[ RUN      ] StochasticTest.test_dropout_rate_stats
[       OK ] StochasticTest.test_dropout_rate_stats
----------------------------------------------------------------------
Ran 21 tests in 13.907s

OK

cgarciae · 2022-09-19T15:52:19Z

@dslisleedh can you create one of more tests under PoolTest using this new argument? You can copy an existing test, rename it, and change it to use the new count_include_pad argument. This would help immensely as these will be checked automatically before merging new code.

…iv_shape for it raises error when there's no batch dimension

dslisleedh · 2022-09-20T11:19:20Z

To @cgarciae

I add some codes to TestPool and here's code I tested.

from absl.testing import absltest, parameterized

from flax import ids
from flax import linen as nn

import jax
from jax import random
from jax import test_util as jtu
from jax.nn import initializers
import jax.numpy as jnp

import numpy as np

# Parse absl flags test_srcdir and test_tmpdir.
jax.config.parse_flags_with_absl()


class PoolTest(parameterized.TestCase):

    def test_pool_custom_reduce(self):
        x = jnp.full((1, 3, 3, 1), 2.)
        mul_reduce = lambda x, y: x * y
        y = nn.pooling.pool(x, 1., mul_reduce, (2, 2), (1, 1), 'VALID')
        np.testing.assert_allclose(y, np.full((1, 2, 2, 1), 2. ** 4))

    @parameterized.parameters(
        {'count_include_pad': True},
        {'count_include_pad': False})
    def test_avg_pool(self, count_include_pad):
        x = jnp.full((1, 3, 3, 1), 2.)
        pool = lambda x: nn.avg_pool(x, (2, 2), count_include_pad=count_include_pad)
        y = pool(x)
        np.testing.assert_allclose(y, np.full((1, 2, 2, 1), 2.))
        y_grad = jax.grad(lambda x: pool(x).sum())(x)
        expected_grad = jnp.array([
            [0.25, 0.5, 0.25],
            [0.5, 1., 0.5],
            [0.25, 0.5, 0.25],
        ]).reshape((1, 3, 3, 1))
        np.testing.assert_allclose(y_grad, expected_grad)

    @parameterized.parameters(
        {'count_include_pad': True},
        {'count_include_pad': False})
    def test_avg_pool_no_batch(self, count_include_pad):
        x = jnp.full((3, 3, 1), 2.)
        pool = lambda x: nn.avg_pool(x, (2, 2), count_include_pad=count_include_pad)
        y = pool(x)
        np.testing.assert_allclose(y, np.full((2, 2, 1), 2.))
        y_grad = jax.grad(lambda x: pool(x).sum())(x)
        expected_grad = jnp.array([
            [0.25, 0.5, 0.25],
            [0.5, 1., 0.5],
            [0.25, 0.5, 0.25],
        ]).reshape((3, 3, 1))
        np.testing.assert_allclose(y_grad, expected_grad)

    def test_max_pool(self):
        x = jnp.arange(9).reshape((1, 3, 3, 1)).astype(jnp.float32)
        pool = lambda x: nn.max_pool(x, (2, 2))
        expected_y = jnp.array([
            [4., 5.],
            [7., 8.],
        ]).reshape((1, 2, 2, 1))
        y = pool(x)
        np.testing.assert_allclose(y, expected_y)
        y_grad = jax.grad(lambda x: pool(x).sum())(x)
        expected_grad = jnp.array([
            [0., 0., 0.],
            [0., 1., 1.],
            [0., 1., 1.],
        ]).reshape((1, 3, 3, 1))
        np.testing.assert_allclose(y_grad, expected_grad)


if __name__ == '__main__':
  absltest.main()

When I tested with the previous code, an error occurred in the non-batch avg_pool, so I corrected the PR. Thanks for telling me to test it for sure.

Below is the test result with the modified code.

(flax_test) idongheon@idongheon-ui-MacBookAir flax % python pool_test.py 
Running tests under Python 3.10.5: /Users/idongheon/miniforge3/envs/flax_test/bin/python
[ RUN      ] PoolTest.test_avg_pool0 (count_include_pad=True)
I0920 20:09:02.610634 4371807552 xla_bridge.py:169] Remote TPU is not linked into jax; skipping remote TPU.
I0920 20:09:02.610964 4371807552 xla_bridge.py:345] Unable to initialize backend 'tpu_driver': Could not initialize backend 'tpu_driver'
I0920 20:09:02.611104 4371807552 xla_bridge.py:345] Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I0920 20:09:02.611212 4371807552 xla_bridge.py:345] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I0920 20:09:02.611675 4371807552 xla_bridge.py:345] Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'
[       OK ] PoolTest.test_avg_pool0 (count_include_pad=True)
[ RUN      ] PoolTest.test_avg_pool1 (count_include_pad=False)
[       OK ] PoolTest.test_avg_pool1 (count_include_pad=False)
[ RUN      ] PoolTest.test_avg_pool_no_batch0 (count_include_pad=True)
[       OK ] PoolTest.test_avg_pool_no_batch0 (count_include_pad=True)
[ RUN      ] PoolTest.test_avg_pool_no_batch1 (count_include_pad=False)
[       OK ] PoolTest.test_avg_pool_no_batch1 (count_include_pad=False)
[ RUN      ] PoolTest.test_max_pool
[       OK ] PoolTest.test_max_pool
[ RUN      ] PoolTest.test_pool_custom_reduce
[       OK ] PoolTest.test_pool_custom_reduce
----------------------------------------------------------------------
Ran 6 tests in 1.586s

OK

Thank you.

cgarciae · 2022-09-20T14:24:53Z

Awesome @dslisleedh! Can you commit changes to the tests?

…iv_shape for it raises error when there's no batch dimension

dslisleedh · 2022-09-20T15:49:40Z

To @cgarciae

Sure :)

cgarciae · 2022-09-20T22:52:34Z

@dslisleedh can you add this test? None of the other tests used padding="SAME" which is the more interesting case.

    @parameterized.parameters(
        {'count_include_pad': True},
        {'count_include_pad': False})
    def test_avg_pool_padding_same(self, count_include_pad):
        x = jnp.array([1.0, 2.0, 3.0, 4.0]).reshape((1, 2, 2, 1))
        pool = lambda x: nn.avg_pool(x, (2, 2), padding="SAME", count_include_pad=count_include_pad)
        y = pool(x)
        if count_include_pad:
          expected_y = jnp.array([10.0 / 4, 6.0 / 4, 7.0 / 4, 4.0 / 4]).reshape((1, 2, 2, 1))
        else:
          expected_y = jnp.array([10.0 / 4, 6.0 / 2, 7.0 / 2, 4.0 / 1]).reshape((1, 2, 2, 1))
        np.testing.assert_allclose(y, expected_y)

dslisleedh · 2022-09-21T13:26:19Z

@cgarciae Oh, I forgot that. Thank you.

and here is result of your code.

(flax_test) idongheon@idongheon-ui-MacBookAir flax % python ./tests/linen/linen_test.py 
Running tests under Python 3.10.5: /Users/idongheon/miniforge3/envs/flax_test/bin/python
[ RUN      ] IdsTest.test_hashable
[       OK ] IdsTest.test_hashable
[ RUN      ] NormalizationTest.test_batch_norm
I0921 22:23:30.726801 4309613888 xla_bridge.py:169] Remote TPU is not linked into jax; skipping remote TPU.
I0921 22:23:30.727137 4309613888 xla_bridge.py:345] Unable to initialize backend 'tpu_driver': Could not initialize backend 'tpu_driver'
I0921 22:23:30.727280 4309613888 xla_bridge.py:345] Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I0921 22:23:30.727386 4309613888 xla_bridge.py:345] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I0921 22:23:30.727836 4309613888 xla_bridge.py:345] Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'
[       OK ] NormalizationTest.test_batch_norm
[ RUN      ] NormalizationTest.test_batch_norm_complex
[       OK ] NormalizationTest.test_batch_norm_complex
[ RUN      ] NormalizationTest.test_batch_norm_multi_init
[       OK ] NormalizationTest.test_batch_norm_multi_init
[ RUN      ] NormalizationTest.test_group_norm
[       OK ] NormalizationTest.test_group_norm
[ RUN      ] NormalizationTest.test_group_norm_raises
[       OK ] NormalizationTest.test_group_norm_raises
[ RUN      ] NormalizationTest.test_layer_norm0 (reduction_axes=-1)
[       OK ] NormalizationTest.test_layer_norm0 (reduction_axes=-1)
[ RUN      ] NormalizationTest.test_layer_norm1 (reduction_axes=1)
[       OK ] NormalizationTest.test_layer_norm1 (reduction_axes=1)
[ RUN      ] NormalizationTest.test_layer_norm2 (reduction_axes=(1, 2))
[       OK ] NormalizationTest.test_layer_norm2 (reduction_axes=(1, 2))
[ RUN      ] PoolTest.test_avg_pool0 (count_include_pad=True)
[       OK ] PoolTest.test_avg_pool0 (count_include_pad=True)
[ RUN      ] PoolTest.test_avg_pool1 (count_include_pad=False)
[       OK ] PoolTest.test_avg_pool1 (count_include_pad=False)
[ RUN      ] PoolTest.test_avg_pool_no_batch0 (count_include_pad=True)
[       OK ] PoolTest.test_avg_pool_no_batch0 (count_include_pad=True)
[ RUN      ] PoolTest.test_avg_pool_no_batch1 (count_include_pad=False)
[       OK ] PoolTest.test_avg_pool_no_batch1 (count_include_pad=False)
[ RUN      ] PoolTest.test_avg_pool_padding_same0 (count_include_pad=True)
[       OK ] PoolTest.test_avg_pool_padding_same0 (count_include_pad=True)
[ RUN      ] PoolTest.test_avg_pool_padding_same1 (count_include_pad=False)
[       OK ] PoolTest.test_avg_pool_padding_same1 (count_include_pad=False)
[ RUN      ] PoolTest.test_max_pool
[       OK ] PoolTest.test_max_pool
[ RUN      ] PoolTest.test_pool_custom_reduce
[       OK ] PoolTest.test_pool_custom_reduce
[ RUN      ] RecurrentTest.test_complex_input_gru
[       OK ] RecurrentTest.test_complex_input_gru
[ RUN      ] RecurrentTest.test_convlstm
[       OK ] RecurrentTest.test_convlstm
[ RUN      ] RecurrentTest.test_gru
[       OK ] RecurrentTest.test_gru
[ RUN      ] RecurrentTest.test_lstm
[       OK ] RecurrentTest.test_lstm
[ RUN      ] RecurrentTest.test_optimized_lstm_cell_matches_regular
/Users/idongheon/miniforge3/envs/flax_test/lib/python3.10/site-packages/jax/test_util.py:44: FutureWarning: jax.test_util.check_eq is deprecated and will soon be removed.
  warnings.warn(f"jax.test_util.{attr} is deprecated and will soon be removed.", FutureWarning)
[       OK ] RecurrentTest.test_optimized_lstm_cell_matches_regular
[ RUN      ] StochasticTest.test_dropout
[       OK ] StochasticTest.test_dropout
[ RUN      ] StochasticTest.test_dropout_rate_limits
[       OK ] StochasticTest.test_dropout_rate_limits
[ RUN      ] StochasticTest.test_dropout_rate_stats
[       OK ] StochasticTest.test_dropout_rate_stats
----------------------------------------------------------------------
Ran 25 tests in 14.023s

OK

cgarciae

This is great @dslisleedh, thanks for going through with this!

dslisleedh added 2 commits September 9, 2022 17:10

Add count_include_pad argument to flax.linen.pooling.avg_pool

9f5ad66

Add count_include_pad argument to flax.linen.pooling.avg_pool

ff97879

cgarciae self-assigned this Sep 13, 2022

cgarciae approved these changes Sep 14, 2022

View reviewed changes

cgarciae mentioned this pull request Sep 14, 2022

Add tests for pooling.py #2464

Closed

Add count_include_pad argument to flax.linen.pooling.avg_pool / Fix d…

8af1605

…iv_shape for it raises error when there's no batch dimension

Add count_include_pad argument to flax.linen.pooling.avg_pool / Fix d…

c73a4b7

…iv_shape for it raises error when there's no batch dimension

Add test_avg_pool_no_batch under test.linen.linen_test.PoolTest

d4fa16d

cgarciae approved these changes Sep 21, 2022

View reviewed changes

cgarciae added the pull ready label Sep 21, 2022

copybara-service Bot merged commit b344022 into google:main Sep 23, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Adding "count_include_pad" argument to flax.linen.pooling.avg_pool#2451

Adding "count_include_pad" argument to flax.linen.pooling.avg_pool#2451
copybara-service[bot] merged 5 commits into
google:mainfrom
dslisleedh:avg_pool_countpad

dslisleedh commented Sep 9, 2022 •

edited

Loading

Uh oh!

codecov-commenter commented Sep 9, 2022 •

edited

Loading

Uh oh!

cgarciae commented Sep 13, 2022

Uh oh!

dslisleedh commented Sep 13, 2022

Uh oh!

cgarciae commented Sep 14, 2022

Uh oh!

dslisleedh commented Sep 14, 2022 •

edited

Loading

Uh oh!

cgarciae commented Sep 19, 2022

Uh oh!

dslisleedh commented Sep 20, 2022

Uh oh!

cgarciae commented Sep 20, 2022

Uh oh!

dslisleedh commented Sep 20, 2022

Uh oh!

cgarciae commented Sep 20, 2022

Uh oh!

dslisleedh commented Sep 21, 2022

Uh oh!

cgarciae left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

dslisleedh commented Sep 9, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Checklist

Uh oh!

codecov-commenter commented Sep 9, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

cgarciae commented Sep 13, 2022

Uh oh!

dslisleedh commented Sep 13, 2022

Uh oh!

cgarciae commented Sep 14, 2022

Uh oh!

dslisleedh commented Sep 14, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cgarciae commented Sep 19, 2022

Uh oh!

dslisleedh commented Sep 20, 2022

Uh oh!

cgarciae commented Sep 20, 2022

Uh oh!

dslisleedh commented Sep 20, 2022

Uh oh!

cgarciae commented Sep 20, 2022

Uh oh!

dslisleedh commented Sep 21, 2022

Uh oh!

cgarciae left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

dslisleedh commented Sep 9, 2022 •

edited

Loading

codecov-commenter commented Sep 9, 2022 •

edited

Loading

dslisleedh commented Sep 14, 2022 •

edited

Loading