A Global Average Pooling Function?

[codymlewis]

Having recently implemented a copy of the ResNet-RS model, I noticed that Flax does not have a function for global average pooling.

So I ended up implementing my own function for use in flax models, just wondering if it may a useful addition to Flax?

[marcvanzee]

I am not very familiar with global average pooling, but doesn't this just reduce to taking the average over all inputs for each feature dimension? Could you please share your code?

[cgarciae]

Hey @codymlewis, global average pooling is just mean over the spatial/temporal dims e.g:

x = jnp.mean(x, axis=(1, 2)) # assuming (batch, height, width, channels)

I don't see a real benefit of adding a GobalAveragePooling Module, currently nn.max_pool and nn.avg_pool are just functions so it would make sense for this operation to also be a function, jnp.mean is already that function so there would need to be some additional motivation to create a dedicated wrapper.

[codymlewis]

Thanks, I agree there. Perhaps I was just overthinking with a more generalized version:

def global_avg_pool(x: Array) -> Array:
	center_axes = tuple(i for i in range(x.ndim) if i != 0 and i != x.ndim - 1)
	return jnp.mean(x, axis=center_axes)

Although that only benefit would be that it would allow people who are not aware that global_avg_pool and jnp.mean are equivalent to use it in their models. But it's probably better design to ensure they understand whats actually happening, as we already have with just jnp.mean.

[marcvanzee]

Thanks for the code snippet! If someone now searches for global average pooling they hopefully will go to this Discussion and then they can use your snippet directly :-)

[cgarciae]

@codymlewis seems good for the typical (batch, *spatial_dims, channels) case, however, most Flax layers such as Dense and Conv support multiple batch dimension (*batch, *spatial_dims, channels).

[codymlewis]

I see, that makes things a little more complex, and probably best off, as previously mentioned, just using normal jnp.mean (e.g. jnp.mean(x, axis=(2, 3)) # for 2D batchs of 3D samples with channels last) or einops (e.g. einops.reduce(x, 'b1 b2 h w c -> b1 b2 c', jnp.mean) # for the same dimensions as the other example). But for the sake of completeness:

def global_avg_pool(x: Array, batch_dims: Sequence[int] = ()) -> Array:
    if not batch_dims:
        batch_dims = (0,)
    center_axes = tuple(i for i in range(x.ndim) if i not in batch_dims and i != x.ndim - 1)
    return jnp.mean(x, axis=center_axes)