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 tf.data.Dataset 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):

  

      github.com
  

  
    Rishit-dagli/Transformer-in-Transformer/blob/main/training.ipynb


    {
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Rishit-dagli/Transformer-in-Transformer/blob/main/training.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FE7KNzPPVrVV"
      },
      "source": [
        "# Image classification with Transformer iN Transformer\n",
        "\n",



  This file has been truncated. show original