DEV Community

Siddhi
Siddhi

Posted on

Exploring Accelerated ML: Training a Neural Network and Understanding TPUs

My First Hands-On ML Experiment: Training a Digit Classifier on Google


What happens when you stop reading about machine learning and actually train a model?

On 24 October, I attended the Accelerated Intelligence Lab organized by GDG. During the workshop, we worked with Google Colab, TensorFlow, and Keras to train a neural network that recognizes handwritten digits.

The instructor provided the implementation and guided us through running it. We followed the notebook, trained the model, and observed its accuracy and validation results.

I wanted to understand more than just whether the model worked. I wanted to understand what was happening between giving it an image and getting a prediction.

Here's what I learned from the experiment.

The experiment: teaching a model to recognize digits

We used the MNIST dataset, which contains images of handwritten digits from 0 to 9. Each image is a 28 × 28 grayscale image.

Our task was straightforward: given an image, predict which digit it represents.

We started by loading the dataset:

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

(x_train, y_train), (x_test, y_test) = (
    keras.datasets.mnist.load_data()
)
Enter fullscreen mode Exit fullscreen mode

The dataset gives us images and their corresponding labels. The training set is used to teach the model, while the test set is reserved for evaluating its performance.

Before training, we normalized the pixel values:

x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0
Enter fullscreen mode Exit fullscreen mode

Why divide by 255? Each pixel originally has an intensity between 0 and 255. Scaling these values to a range between 0 and 1 makes the input values more manageable for the neural network.

It's a small preprocessing step, but it illustrates something important: preparing the data is part of building a machine learning model.

Building the neural network

We used a simple sequential neural network:

model = keras.Sequential([
    keras.Input(shape=(28, 28)),
    layers.Flatten(),
    layers.Dense(128, activation="relu"),
    layers.Dense(10, activation="softmax")
])
Enter fullscreen mode Exit fullscreen mode

This architecture has three main layers after the input.

The Flatten layer converts each 28 × 28 image into a one-dimensional array of 784 values.

The hidden Dense layer contains 128 neurons and uses ReLU activation. It learns patterns from the input pixels.

The final layer has 10 output units, one for each possible digit. Softmax converts the output scores into values that can be interpreted as class probabilities.

The model doesn't start knowing what a handwritten 7 looks like. It learns patterns by adjusting its parameters during training.

That's the part I find interesting: we provide examples and their labels, and the model gradually adjusts itself to make better predictions.

Training the model

After defining the architecture, we configured it for training:

model.compile(
    optimizer="adam",
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"]
)
Enter fullscreen mode Exit fullscreen mode

We then trained it for five epochs, using a batch size of 128 and reserving 10% of the training data for validation:

history = model.fit(
    x_train,
    y_train,
    epochs=5,
    batch_size=128,
    validation_split=0.1
)
Enter fullscreen mode Exit fullscreen mode

An epoch represents one pass through the training dataset. The batch size determines how many training examples are processed together before the model updates its parameters.

During training, we observed two metrics: training accuracy and validation accuracy.

Training accuracy measures performance on the training examples. Validation accuracy measures performance on the held-out validation portion.

Our accuracy plot showed both curves increasing over the five epochs. By the end, both were approaching 98%.

That was a useful result to observe, but accuracy alone doesn't tell the whole story. We also need to evaluate the trained model on the separate test set to estimate how well it performs on data it wasn't trained on.

Testing the model's predictions

After training, we evaluated the model on the test dataset:

test_loss, test_accuracy = model.evaluate(
    x_test, y_test, verbose=0
)

print(f"Test accuracy: {test_accuracy * 100:.2f}%")
Enter fullscreen mode Exit fullscreen mode

We also used the model to predict a digit from an individual image:

index = 0

prediction = model.predict(
    x_test[index:index+1], verbose=0
)

predicted_digit = prediction.argmax()

print("Predicted:", predicted_digit)
print("Actual:", y_test[index])
Enter fullscreen mode Exit fullscreen mode

The model produces scores for the 10 possible digits. Taking the index of the highest score gives us its predicted class.

We also visualized predictions for 10 test images, comparing the predicted digit with the actual label.

I liked this part because it made the model's output tangible. Instead of looking only at a number representing accuracy, we could inspect individual predictions and see where the model was correct or incorrect.

Exploring TPUs and accelerated machine learning

The workshop also introduced us to TPUs and the broader ecosystem of tools used for accelerated machine learning.

Our Colab notebook showed a TPU runtime, and we explored device availability using TensorFlow. This introduced us to the connection between the code we write and the hardware available to execute it.

A TPU is a specialized processor designed to accelerate machine learning computations. GPUs can also accelerate these workloads, but the best choice depends on the model, hardware, and software configuration.

One important distinction I learned is that selecting a TPU runtime is not the same as proving that the model trained on a TPU.

In TensorFlow, proper TPU execution typically involves connecting to the TPU, creating a TPUStrategy, and constructing the model within the strategy's scope. Our notebook's runtime and device checks are useful starting points, but the exact training configuration needs to be inspected before drawing conclusions about accelerator usage.

This gave me a reason to explore the implementation more deeply rather than treating hardware acceleration as just another setting in Coll

ab.

What I'm taking away from this workshop

This was a guided implementation rather than a model I developed independently, but running the notebook gave me a practical starting point.

I got to see how dataset loading, preprocessing, model architecture, training, validation, and prediction fit together in one workflow.

It also raised questions I want to investigate further: how different architectures affect accuracy, how to interpret misclassified examples, and whether changing the hardware makes a measurable difference to training time.

My next step is to revisit the code and experiment with it independently, rather than simply rerunning the same implementation.

For me, that is the difference between completing a workshop exercise and turning it into a learning project.

Thanks to GDG for organizing the session and giving us an opportunity to explore these tools hands-on.

MachineLearning #TensorFlow #GoogleColab #AI #GDG

Top comments (0)