How to Perform Image Classification with TensorFlow on Ubuntu 24.04 GPU Server

Image classification is one of the most common tasks in deep learning. It allows you to train a model that can recognize and categorize images into predefined classes, for example, identifying whether a picture shows a cat, dog, or car. TensorFlow, a popular open-source deep learning framework, makes it easier to build, train, and deploy image classification models.

In this tutorial, you will learn how to perform image classification on an Ubuntu 24.04 GPU server using TensorFlow.

Prerequisites

• An Ubuntu 24.04 server with an NVIDIA GPU.
• A non-root user or a user with sudo privileges.
• NVIDIA drivers are installed on your server.

Step 1 – Setting Up the Python Environment

Before building and training your image classification model, you need to set up a dedicated Python environment with TensorFlow installed.

1. Update your system and install Python, pip, venv, and Git.

apt update -y
apt install -y python3 python3-venv python3-pip git

2. Create a new virtual environment for TensorFlow.

python3 -m venv tf-gpu-env
source tf-gpu-env/bin/activate

3. Upgrade pip and install TensorFlow.

pip install --upgrade pip
pip install tensorflow

4. Confirm TensorFlow detects your GPU.

python3 -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

Output.

[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]

This means TensorFlow is ready to use your GPU for faster training.

Step 2 – Preparing the Dataset

Now that your environment is ready, the next step is to load and preprocess the dataset you’ll use for training the image classification model. In this tutorial, we’ll work with the CIFAR-10 dataset, a benchmark dataset for image classification tasks.

1. Create a dataset script.

nano dataset_prep.py

Add the following code.

# save as dataset_prep.py
import tensorflow as tf

# Load CIFAR-10 dataset
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()

# Normalize images
x_train, x_test = x_train / 255.0, x_test / 255.0

print("Training data shape:", x_train.shape)
print("Test data shape:", x_test.shape)

2. Execute the script to verify that the dataset is loaded correctly.

python3 dataset_prep.py

Output.

170498071/170498071 ━━━━━━━━━━━━━━━━━━━━ 7s 0us/step   
Training data shape: (50000, 32, 32, 3)
Test data shape: (10000, 32, 32, 3)

The above output tells that:

• Training set has 50,000 images.
• Test set has 10,000 images.
• Each image is 32×32 pixels with 3 color channels (RGB).

By normalizing the images, we make training more efficient and stable. The dataset is now ready for model building.

Step 3 – Building the CNN Model

To classify images from the CIFAR-10 dataset, we’ll use a Convolutional Neural Network (CNN). CNNs are well-suited for image tasks because they can automatically detect patterns such as edges, textures, and shapes.

1. Create the model script.

nano cnn_model.py

Add the following code.

# save as cnn_model.py
import tensorflow as tf
from tensorflow.keras import layers, models

def build_model():
    model = models.Sequential([
        layers.Conv2D(32, (3,3), activation='relu', input_shape=(32,32,3)),
        layers.MaxPooling2D((2,2)),
        layers.Conv2D(64, (3,3), activation='relu'),
        layers.MaxPooling2D((2,2)),
        layers.Conv2D(64, (3,3), activation='relu'),
        layers.Flatten(),
        layers.Dense(64, activation='relu'),
        layers.Dense(10, activation='softmax')
    ])
    return model

if __name__ == "__main__":
    model = build_model()
    model.summary()

2. Run the script to see the model’s architecture.

python3 cnn_model.py

You should see the following model summary:

Model: "sequential"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┓
┃ Layer (type)                         ┃ Output Shape                ┃         Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━┩
│ conv2d (Conv2D)                      │ (None, 30, 30, 32)          │             896 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ max_pooling2d (MaxPooling2D)         │ (None, 15, 15, 32)          │               0 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ conv2d_1 (Conv2D)                    │ (None, 13, 13, 64)          │          18,496 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ max_pooling2d_1 (MaxPooling2D)       │ (None, 6, 6, 64)            │               0 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ conv2d_2 (Conv2D)                    │ (None, 4, 4, 64)            │          36,928 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ flatten (Flatten)                    │ (None, 1024)                │               0 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ dense (Dense)                        │ (None, 64)                  │          65,600 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ dense_1 (Dense)                      │ (None, 10)                  │             650 │
└──────────────────────────────────────┴─────────────────────────────┴─────────────────┘
 Total params: 122,570 (478.79 KB)
 Trainable params: 122,570 (478.79 KB)
 Non-trainable params: 0 (0.00 B)

Step 4 – Training the Model

With the dataset prepared and CNN model defined, the next step is to train the model on the CIFAR-10 dataset. Training allows the CNN to learn patterns in the images and improve its ability to classify them correctly.

1. Create a training script.

nano train_model.py

Add the following code.

# train_model.py
import tensorflow as tf
import os
from tensorflow.keras import datasets
from cnn_model import build_model

(x_train, y_train), (x_test, y_test) = datasets.cifar10.load_data()
x_train, x_test = x_train/255.0, x_test/255.0

model = build_model()

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

# Train model
history = model.fit(x_train, y_train, epochs=10,
                    validation_data=(x_test, y_test))

# ✅ Ensure save directory exists
os.makedirs("saved_model", exist_ok=True)

# Save full model and weights
model.save("saved_model/cnn.keras")           # preferred modern format
model.save_weights("saved_model/cnn.weights.h5")  # backup (weights only)

2. Run the training script.

python3 train_model.py

During training, you’ll see output logs for each epoch.

1563/1563 ━━━━━━━━━━━━━━━━━━━━ 14s 5ms/step - accuracy: 0.4384 - loss: 1.5377 - val_accuracy: 0.5465 - val_loss: 1.2606
Epoch 2/10
1563/1563 ━━━━━━━━━━━━━━━━━━━━ 3s 2ms/step - accuracy: 0.5848 - loss: 1.1710 - val_accuracy: 0.6242 - val_loss: 1.0618
Epoch 3/10
1563/1563 ━━━━━━━━━━━━━━━━━━━━ 3s 2ms/step - accuracy: 0.6378 - loss: 1.0266 - val_accuracy: 0.6493 - val_loss: 0.9867
Epoch 4/10
1563/1563 ━━━━━━━━━━━━━━━━━━━━ 3s 2ms/step - accuracy: 0.6723 - loss: 0.9303 - val_accuracy: 0.6738 - val_loss: 0.9508
Epoch 5/10
1563/1563 ━━━━━━━━━━━━━━━━━━━━ 3s 2ms/step - accuracy: 0.6951 - loss: 0.8659 - val_accuracy: 0.6778 - val_loss: 0.9201
Epoch 6/10
1563/1563 ━━━━━━━━━━━━━━━━━━━━ 3s 2ms/step - accuracy: 0.7188 - loss: 0.8065 - val_accuracy: 0.6868 - val_loss: 0.9126
Epoch 7/10
1563/1563 ━━━━━━━━━━━━━━━━━━━━ 3s 2ms/step - accuracy: 0.7342 - loss: 0.7576 - val_accuracy: 0.6918 - val_loss: 0.8812
Epoch 8/10
1563/1563 ━━━━━━━━━━━━━━━━━━━━ 3s 2ms/step - accuracy: 0.7501 - loss: 0.7130 - val_accuracy: 0.7058 - val_loss: 0.8600
Epoch 9/10
1563/1563 ━━━━━━━━━━━━━━━━━━━━ 3s 2ms/step - accuracy: 0.7637 - loss: 0.6772 - val_accuracy: 0.7079 - val_loss: 0.8726
Epoch 10/10
1563/1563 ━━━━━━━━━━━━━━━━━━━━ 3s 2ms/step - accuracy: 0.7759 - loss: 0.6436 - val_accuracy: 0.6847 - val_loss: 0.9227

Here you can see:

• Training accuracy improving with each epoch.
• Validation accuracy reaching around ~68%.
• The trained model and weights are saved in the saved_model/ directory.

At this stage, you have a trained CNN ready for evaluation.

Step 5 – Evaluating the Model

After training, the next step is to evaluate the model’s performance on the test dataset. This helps you understand how well the model generalizes to unseen data.

1. Create the evaluation script.

nano evaluate.py

Add the following code.

# evaluate.py
import os
import tensorflow as tf
from tensorflow.keras import datasets
from cnn_model import build_model

# Reload CIFAR-10 dataset
(_, _), (x_test, y_test) = datasets.cifar10.load_data()
x_test = x_test / 255.0

model = None

# ✅ Try full model first
if os.path.exists("saved_model/cnn.keras"):
    print("Loading full model from saved_model/cnn.keras")
    model = tf.keras.models.load_model("saved_model/cnn.keras")

# ✅ Fallback to weights-only
elif os.path.exists("saved_model/cnn.weights.h5"):
    print("Loading model architecture + weights from saved_model/cnn.weights.h5")
    model = build_model()
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    model.load_weights("saved_model/cnn.weights.h5")

else:
    raise FileNotFoundError("No saved model found! Please run train_model.py first.")

# Evaluate model
loss, acc = model.evaluate(x_test, y_test, verbose=2)
print(f"Test accuracy: {acc*100:.2f}%")

2. Run the evaluation script.

python3 evaluate.py

You should see something like this:

313/313 - 6s - 18ms/step - accuracy: 0.6847 - loss: 0.9227
Test accuracy: 68.47%

This indicates that the trained CNN correctly classified approximately 68% of the test images. While this isn’t state-of-the-art, it’s a solid baseline for CIFAR-10 using a relatively simple CNN.

Step 6 – Making Predictions

Now that the model has been trained and evaluated, let’s use it to make predictions on new data. This step demonstrates how to reload a saved model and test it against unseen images.

1. Create a prediction script.

nano predict.py

Add the following code.

# predict.py
import os
import numpy as np
import tensorflow as tf
from tensorflow.keras import datasets
from cnn_model import build_model

# Reload CIFAR-10 test dataset
(_, _), (x_test, y_test) = datasets.cifar10.load_data()
x_test = x_test / 255.0

model = None

# ✅ Try full model first
if os.path.exists("saved_model/cnn.keras"):
    print("Loading full model from saved_model/cnn.keras")
    model = tf.keras.models.load_model("saved_model/cnn.keras")

# ✅ Fallback to weights-only
elif os.path.exists("saved_model/cnn.weights.h5"):
    print("Loading model architecture + weights from saved_model/cnn.weights.h5")
    model = build_model()
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    model.load_weights("saved_model/cnn.weights.h5")

else:
    raise FileNotFoundError("No saved model found! Please run train_model.py first.")

# ✅ Make predictions on first 5 test samples
predictions = model.predict(x_test[:5])

print("Predicted labels:", np.argmax(predictions, axis=1))
print("True labels     :", y_test[:5].flatten())

2. Run the script.

python3 predict.py

You should see output similar to this:

Predicted labels: [3 1 8 0 4]
True labels     : [3 8 8 0 6]

Here, the model predicted labels for the first 5 test images.

• It correctly predicted 3 (cat) and 0 (airplane).
• It misclassified some others, which is common with a relatively simple CNN trained for only 10 epochs.

Conclusion

In this tutorial, you built an image classification pipeline with TensorFlow on Ubuntu 24.04 GPU server. You prepared the CIFAR-10 dataset, created a CNN model, trained it, evaluated its accuracy, and made predictions on test images. The model reached around 68% accuracy, which is a solid baseline for a simple CNN trained for 10 epochs.