Translating a basic MNIST tensorflow network into an MLJFlux builder

On slack I have been asked:

How does one translate the following into MLJFlux. The code is from this tensorflow tutorial on the MNIST dataset for beginners.

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
1 Like

You can define an equivalent MLJFlux chain builder (that works with images of any size) with:

import MLJFlux
using Flux

builder = MLJFlux.@builder Chain(
    Dense(prod(n_in)*n_channels => 128, relu),
    Dense(128, 10),

Just change the builder in this notebook to the one above to use with MNIST images.

Note that ImageClassifier automatically adds a softmax at the end, to get probabilities. If you don’t want that final layer, set finaliser=identity in the ImageClassifier constructor.