The Beginner‘s Guide to Training a Classification Model with TensorFlow

Machine learning powers many intelligent applications we use every day, from spam filters to product recommendations. Classification is one of the core tasks in machine learning, used for everything from diagnosing diseases to detecting credit card fraud. In this tutorial, you‘ll learn how to build your own classification model using the popular TensorFlow library.

We‘ll walk through the process step-by-step, from gathering data to evaluating the trained model. By the end, you‘ll be equipped with the skills to tackle your own classification problems. Let‘s dive in!

What is Classification?

In machine learning, classification is the task of predicting which category an input belongs to, based on labeled training data. A classic example is classifying emails as "spam" or "not spam". The model learns to recognize patterns from example emails and their labels, then uses this knowledge to classify new emails it hasn‘t seen before.

Some other common applications of classification include:

• Medical diagnosis: Identifying diseases based on patient symptoms and test results
• Sentiment analysis: Determining if a movie review or social media post is positive, negative, or neutral
• Fraud detection: Flagging suspicious transactions as potential fraud

The classification model takes in an input, extracts informative features, and outputs the predicted probability for each possible class label. The class with the highest predicted probability is chosen as the final prediction.

Why TensorFlow?

TensorFlow is a powerful open-source library for machine learning developed by Google. It allows you to build models by connecting mathematical operations into computational graphs. Don‘t let that intimidate you though – TensorFlow also provides a more beginner-friendly interface called Keras that abstracts away a lot of the low-level complexity.

Some key advantages of TensorFlow include:

• Flexibility to build all kinds of models, from simple linear classifiers to complex neural networks
• Ability to train on CPUs, GPUs, or TPUs for faster performance
• Portability to deploy trained models on servers, desktops, mobile devices, and even embedded systems
• Strong community support and a wealth of documentation and tutorials

With TensorFlow 2.0 and above, eager execution is enabled by default which allows for a more intuitive imperative programming style. This means you can inspect tensors and debug your code line by line like you would with NumPy.

Gathering and Preparing Data

Having a suitable dataset is crucial for training an effective classification model. The dataset should be representative of the problem you‘re trying to solve and have sufficient examples for the model to learn from. It‘s also important that the classes are well-defined and the labels are accurate.

For this tutorial, we‘ll use the classic Iris flower dataset. This dataset contains measurements for 150 iris flowers from three different species. Our goal will be to train a model to predict the species based on the measurement features.

First, we‘ll load the dataset using the Pandas library:

import pandas as pd

df = pd.read_csv(‘https://archive.ics.uci.edu/ml/machine-learning-databases/iris/iris.data‘, header=None)
df.columns = [‘sepal_length‘, ‘sepal_width‘, ‘petal_length‘, ‘petal_width‘, ‘species‘]
df.head()

Output:

	sepal_length	sepal_width	petal_length	petal_width	species
0	5.1	3.5	1.4	0.2	Iris-setosa
1	4.9	3.0	1.4	0.2	Iris-setosa
2	4.7	3.2	1.3	0.2	Iris-setosa
3	4.6	3.1	1.5	0.2	Iris-setosa
4	5.0	3.6	1.4	0.2	Iris-setosa

We can see there are 4 numeric features and 1 categorical feature for the species. Let‘s check if there are any missing values:

df.isnull().sum()

Output:

sepal_length    0
sepal_width     0
petal_length    0
petal_width     0
species         0
dtype: int64

Great, the data is clean with no missing values! Next, we‘ll convert the text species labels to numeric values. Sklearn‘s LabelEncoder can do this for us:

from sklearn.preprocessing import LabelEncoder

encoder = LabelEncoder()
df[‘species‘] = encoder.fit_transform(df[‘species‘])
df.head()

Output:

	sepal_length	sepal_width	petal_length	petal_width	species
0	5.1	3.5	1.4	0.2	0
1	4.9	3.0	1.4	0.2	0
2	4.7	3.2	1.3	0.2	0
3	4.6	3.1	1.5	0.2	0
4	5.0	3.6	1.4	0.2	0

The final step is to split our data into training and test sets:

from sklearn.model_selection import train_test_split

X = df.drop(‘species‘, axis=1) 
y = df[‘species‘]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

We‘ve used 80% of the data for training and held out 20% for testing. Setting the random state ensures we get the same train-test split each time.

Building the Model

Now we‘re ready to build our neural network classifier using TensorFlow and Keras. Since this is a multiclass classification problem with 3 possible species, we‘ll use a softmax output layer to get the predicted probabilities for each class.

Here‘s how we can define the model architecture:

from tensorflow import keras
from tensorflow.keras import layers

model = keras.Sequential([
  layers.Dense(64, activation=‘relu‘, input_shape=[4]),
  layers.Dense(64, activation=‘relu‘),
  layers.Dense(3, activation=‘softmax‘)
])

model.compile(optimizer=‘adam‘,
              loss=‘sparse_categorical_crossentropy‘,
              metrics=[‘accuracy‘])

model.summary()

Output:

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense (Dense)                (None, 64)                320       
_________________________________________________________________
dense_1 (Dense)              (None, 64)                4160      
_________________________________________________________________
dense_2 (Dense)              (None, 3)                 195       
=================================================================
Total params: 4,675
Trainable params: 4,675
Non-trainable params: 0
_________________________________________________________________

We‘ve defined a sequential model with 3 dense layers:

• The first layer has 64 neurons and uses the ReLU activation function. It expects inputs of shape (4,) matching our 4 features.
• The second hidden layer also has 64 neurons and ReLU activation.
• The final output layer has 3 neurons (one per class) and uses softmax to output a probability distribution.

The model is compiled with the Adam optimizer, sparse categorical cross entropy loss, and accuracy metric. Sparse categorical cross entropy is used when the labels are integers, as opposed to one-hot encoded vectors.

Training the Model

With our model architecture defined, we can now fit it to the training data:

history = model.fit(X_train, y_train, epochs=50, validation_split=0.2)

We‘ve trained for 50 epochs, using 20% of the training data for validation. This allows us to monitor the model‘s performance on unseen data during training.

Let‘s plot the learning curves to see how the model‘s loss and accuracy evolved over time:

import matplotlib.pyplot as plt

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))

ax1.plot(history.history[‘loss‘], label=‘Train‘)
ax1.plot(history.history[‘val_loss‘], label=‘Validation‘)
ax1.set_title(‘Model loss‘)
ax1.set_ylabel(‘Loss‘)
ax1.set_xlabel(‘Epoch‘)
ax1.legend()

ax2.plot(history.history[‘accuracy‘], label=‘Train‘) 
ax2.plot(history.history[‘val_accuracy‘], label=‘Validation‘)
ax2.set_title(‘Model accuracy‘)
ax2.set_ylabel(‘Accuracy‘)
ax2.set_xlabel(‘Epoch‘)
ax2.legend();

Output:

We can see the training and validation loss decreased over epochs, while the accuracy increased, indicating the model learned to fit the data. There are no major gaps between the training and validation curves, which is a good sign the model hasn‘t overfit.

Evaluating Model Performance

Finally, we‘ll evaluate our trained model on the held-out test set to assess its performance on completely unseen data:

test_loss, test_acc = model.evaluate(X_test, y_test)
print(f"Test accuracy: {test_acc:.3f}")

Output:

Test accuracy: 0.967

Our model reaches an impressive 96.7% classification accuracy on the test set!

To understand what kind of errors it‘s making, we can plot a confusion matrix:

from sklearn.metrics import confusion_matrix
import seaborn as sns

pred_probs = model.predict(X_test)
preds = pred_probs.argmax(axis=1)

conf_mat = confusion_matrix(y_test, preds)

plt.figure(figsize=(8,6))
sns.heatmap(conf_mat, annot=True, fmt=‘d‘, cmap=‘Blues‘, cbar=False, 
            xticklabels=encoder.classes_, yticklabels=encoder.classes_);
plt.xlabel(‘Predicted label‘)
plt.ylabel(‘True label‘)
plt.title(‘Confusion Matrix‘);

Output:

Each row corresponds to the true species, while columns are the model‘s predicted species. The diagonal entries represent correct classifications. We can see the model performs well across all 3 classes, with only a couple misclassifications between Virginica and Versicolor.

Saving and Using the Model

Now that we have a trained model we‘re satisfied with, we can save it to use later:

model.save(‘iris_classifier.h5‘)

The model architecture, weights, and optimizer state are saved in the HDF5 file. We can load this model anytime to make predictions on new data:

loaded_model = keras.models.load_model(‘iris_classifier.h5‘)

new_data = [[5.1, 3.5, 1.4, 0.2],
            [6.7, 3.1, 5.6, 2.4],
            [6.4, 2.8, 5.6, 2.1]]

preds = loaded_model.predict(new_data).argmax(axis=1)
print(f"Predicted classes: {preds}")
print(f"Predicted species: {encoder.inverse_transform(preds)}")

Output:

Predicted classes: [0 2 2]
Predicted species: [‘Iris-setosa‘ ‘Iris-virginica‘ ‘Iris-virginica‘]

Where to Go From Here

Congrats, you‘ve just trained your first classification model with TensorFlow and Keras! Some next steps you could explore:

• Experiment with different model architectures
• Tune the hyperparameters using techniques like random search or Bayesian optimization
• Apply regularization to further reduce overfitting
• Explore other classification algorithms like support vector machines, random forests, etc
• Deploy your trained model as a microservice or integrate into an application

This tutorial provides a foundation you can build upon for more complex, real-world classification problems. Don‘t hesitate to dive deeper into TensorFlow and discover all the possibilities it enables!