Converting channel_last data to channel_first

I was writing a tutorial for one of my paper implementations quite similar to the Image classification tutorial and the model I use expects data in the channels_first format. Of course, while using we can simply use tf.transpose.

I was wondering if there was any built-in way to load data with tf.keras.utils.image_dataset_from_directory in the channels_first format or a layer to do so?

Right now I do so in my model adding a lambda layer to do so:

model = tf.keras.Sequential([tf.keras.layers.Lambda(lambda x: tf.einsum("...ijk->...kij", x)),

Here is my complete Colab notebook if required (created in a similar fashion to the classification guide on TensorFlow website):

Another option is to use the tf.keras.layers.Permute

Link: tf.keras.layers.Permute  |  TensorFlow Core v2.8.0


In the end this API is just going to dataset map over load_image that it is a multiple call over tf.image.decode_image so you don’t have any parameter to do it internally.


Got it, thanks so much @ariG23498 and @Bhack ! Also surprisingly enough (I haven’t checked the source code for Permute layer yet and I will get to doing that) while my experiments I noticed the exact same operation with tf.keras.layers.Permute is faster than tf.einsum.

1 Like

permute is a tf.transpose

1 Like