Does flax/jax care about image channel order?

[davidshen84]

Hi,

Most ML libraries prefer images to be pre-processed to be channel first. But I could not find a reference about how flax/jax prefer it.

Does flax take care of it automatically? (I don't think it is possible)

[chiamp]

For the Conv layer, it's channels-last convention

[epignatelli]

That's the correct answer, the one marked as "answer" is incorrect at the date of this comment.
See the docs

[davidshen84]

Thanks. I noticed the underlying implementation uses https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv_general_dilated.html, which defaults channel-first. Is there any reason we changed the convention in flax?

…

On Wed, 29 Nov 2023 at 06:57, Marcus Chiam ***@***.***> wrote: For the Conv layer, it's *channels-last* convention, see: https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.Conv.html#flax.linen.Conv.__call __ — Reply to this email directly, view it on GitHub <#3501 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAAQBTOHIAC3GN725HWLKX3YGY6ZRAVCNFSM6AAAAAA7W42RTOVHI2DSMVQWIX3LMV43SRDJONRXK43TNFXW4Q3PNVWWK3TUHM3TMOJWGY2DE> . You are receiving this because you authored the thread.Message ID: ***@***.***>

[chiamp]

perhaps @marcvanzee would know

[ChandraPrakash-Bathula]

Flax and JAX themselves do not have a strict preference for image channel order, and they don't inherently assume a specific order for image channels (either channel-first or channel-last). However, the choice of channel order is often determined by the underlying libraries and frameworks used in conjunction with Flax or JAX.

[epignatelli]

Where do you specify which convention to choose?

[ChandraPrakash-Bathula]

This might be helpful : https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.Conv.__call

Thank you @epignatelli for the reference.

[epignatelli]

In the docs I read:

inputs – input data with dimensions (*batch_dims, spatial_dims…, features). This is the channels-last convention, i.e. NHWC for a 2d convolution and NDHWC for a 3D convolution.