Classify iris images using Python’s Flax CNN

In my last post I classified images as being either horses or humans using a Flax Convolutional Neural Network (CNN). the link to my last blog post can be found here:- https://medium.com/@tracyrenee61/determine-whether-an-image-is-a-horse-or-a-human-using-a-cnn-made-with-flax-a6f2cb6a5696

Even though the dataset that I used was a binary classification and in theory the predictions should have been solved using a sigmoid activation function, my studies revealed that when it comes to CNNs, a softmax activation function is used whether the dataset is a binary classification or multinomial classification model.

In this blog post I intend to review the code of a multinomial classification model and have selected the iris image dataset for this purpose. The Kaggle dataset can be found here:- https://www.kaggle.com/datasets/jeffheaton/iris-computer-vision/data

The model that has been used to classify the three classes of iris is a CNN written by Flax. A CNN consists of an input layer, hidden layers and an output layer. In a convolutional neural network, the hidden layers include one or more layers that perform convolutions. Typically this includes a layer that performs a dot product of the convolution kernel with the layer’s input matrix. This product is usually the Frobenius inner product, and its activation function is…