Image recognition has become a cornerstone of modern AI applications, from autonomous vehicles to medical imaging. While pre‑trained models are handy, building a custom pipeline lets you tailor data preprocessing, model architecture, and deployment to your exact needs. In this guide we’ll walk through every stage of constructing a TensorFlow‑based image recognition system, from gathering raw images to serving the model via TensorFlow Serving. The steps are written for developers comfortable with Python, Linux command‑line tools, and basic machine‑learning concepts.
What You’ll Need
- Python 3.9+ (preferably via
pyenvorconda) - TensorFlow 2.12 or newer
- Git and a GitHub account for version control
- Docker (for containerised serving)
- GPU‑enabled workstation or cloud instance (optional but recommended)
- Dataset of labeled images (e.g., a folder structure
data/train/class_x/...) - Basic Linux utilities:
wget,tar,ffmpeg(if handling video frames)
Step 1: Set Up a Clean Development Environment
Creating an isolated environment prevents package conflicts and makes reproducibility easier. We’ll use conda, but venv works as well.
# Create a new environment named tf-image
conda create -n tf-image python=3.10 -y
conda activate tf-image
# Install TensorFlow with GPU support if you have CUDA installed
pip install "tensorflow[and-cuda]" # or simply pip install tensorflow for CPU only
# Additional utilities
pip install numpy pandas matplotlib scikit-learn pillow tqdm
Verify the installation:
python -c "import tensorflow as tf; print(tf.__version__)"
If you see a version number ≥ 2.12, you’re ready to go.
Step 2: Acquire and Organise Your Dataset
A well‑structured dataset is the foundation of any successful model. For this guide we’ll use the publicly available Fashion‑MNIST as a placeholder, but you can replace it with any custom collection.
# Create a data directory
mkdir -p data/raw && cd data/raw
# Download the archive (example using wget)
wget https://github.com/zalandoresearch/fashion-mnist/archive/refs/heads/master.zip -O fashion-mnist.zip
unzip fashion-mnist.zip
# Re‑organise into TensorFlow‑friendly layout
mkdir -p ../../train/class_0 ../../train/class_1
# (Copy images into appropriate class folders)
TensorFlow’s image_dataset_from_directory expects a directory per class. Ensure each class folder contains only images of that class and that file names are unique.
Step 3: Preprocess Images and Build a Data Pipeline
Efficient data loading is crucial for GPU utilisation. We’ll use the tf.data API to create a pipeline that reads, decodes, resizes, and normalises images on the fly.
import tensorflow as tf
import pathlib
DATA_DIR = pathlib.Path('../data')
BATCH_SIZE = 32
IMG_SIZE = (224, 224) # MobileNetV2 default size
train_ds = tf.keras.utils.image_dataset_from_directory(
DATA_DIR / 'train',
validation_split=0.2,
subset='training',
seed=42,
image_size=IMG_SIZE,
batch_size=BATCH_SIZE)
val_ds = tf.keras.utils.image_dataset_from_directory(
DATA_DIR / 'train',
validation_split=0.2,
subset='validation',
seed=42,
image_size=IMG_SIZE,
batch_size=BATCH_SIZE)
# Normalise pixel values to [0,1]
normalisation_layer = tf.keras.layers.Rescaling(1./255)
train_ds = train_ds.map(lambda x, y: (normalisation_layer(x), y))
val_ds = val_ds.map(lambda x, y: (normalisation_layer(x), y))
# Prefetch for performance
AUTOTUNE = tf.data.AUTOTUNE
train_ds = train_ds.cache().shuffle(1000).prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)
Common mistake: forgetting to set seed consistently for training/validation splits, which leads to overlapping images between sets.
Step 4: Choose a Base Model and Fine‑Tune
Starting from a pre‑trained backbone (e.g., MobileNetV2) accelerates convergence and often yields higher accuracy than training from scratch.
# Load MobileNetV2 without the top classification layer
base_model = tf.keras.applications.MobileNetV2(input_shape=IMG_SIZE + (3,),
include_top=False,
weights='imagenet')
base_model.trainable = False # Freeze base for initial training
# Add custom classification head
global_average_layer = tf.keras.layers.GlobalAveragePooling2D()
prediction_layer = tf.keras.layers.Dense(train_ds.cardinality().numpy(), activation='softmax')
model = tf.keras.Sequential([
base_model,
global_average_layer,
tf.keras.layers.Dropout(0.2),
prediction_layer
])
model.compile(optimizer=tf.keras.optimizers.Adam(),
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
model.summary()
Tip: Use model.layers[-1].output_shape[-1] to automatically match the number of classes.
Step 5: Train, Evaluate, and Fine‑Tune the Model
We’ll first train the top layers while keeping the backbone frozen, then unfreeze some layers for fine‑tuning.
# Initial training (5‑7 epochs is usually enough for a small dataset)
initial_epochs = 5
history = model.fit(train_ds,
validation_data=val_ds,
epochs=initial_epochs)
# Unfreeze the last 50 layers of the base model for fine‑tuning
base_model.trainable = True
for layer in base_model.layers[:-50]:
layer.trainable = False
model.compile(optimizer=tf.keras.optimizers.Adam(1e-5), # lower LR for fine‑tuning
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
fine_tune_epochs = 5
total_epochs = initial_epochs + fine_tune_epochs
history_fine = model.fit(train_ds,
validation_data=val_ds,
epochs=total_epochs,
initial_epoch=history.epoch[-1])
Common mistake: using a learning rate that’s too high during fine‑tuning, which can quickly erase the pre‑trained weights.
Step 6: Export the Model for Production
TensorFlow SavedModel format works seamlessly with TensorFlow Serving and TensorFlow Lite. We’ll also create a versioned directory for easy rollout.
import datetime, os
export_path = f'./saved_model/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}'
model.save(export_path, save_format='tf')
print(f'Model exported to {export_path}')
To test the SavedModel locally:
loaded = tf.keras.models.load_model(export_path)
# Run a single inference on a batch from the validation set
for images, labels in val_ds.take(1):
preds = loaded.predict(images)
print('Predictions shape:', preds.shape)
break
Step 7: Containerise with Docker and Serve via TensorFlow Serving
Docker isolates the serving environment and makes scaling straightforward.
# Dockerfile (place in the same folder as the SavedModel)
FROM tensorflow/serving:2.12.0
COPY ./saved_model /models/custom_image_recognition/1
ENV MODEL_NAME=custom_image_recognition
# Build the image
docker build -t custom_tf_serving .
# Run the container (expose port 8501 for REST API)
docker run -p 8501:8501 --name tf_serving_custom custom_tf_serving
Verify the REST endpoint:
curl -X POST http://localhost:8501/v1/models/custom_image_recognition:predict
-d '{"instances": [$(python -c "import numpy as np, tensorflow as tf, json; img = tf.keras.preprocessing.image.load_img('data/train/class_0/example.jpg', target_size=(224,224)); arr = tf.keras.preprocessing.image.img_to_array(img)/255.0; print(json.dumps(arr.tolist()))")]}'
Common mistake: forgetting to set the model version folder (1) inside /models/…, which causes TensorFlow Serving to return a 404.
Common Mistakes to Avoid
1. Mixing up class indices. The order of folders determines label indices. Always print train_ds.class_names after loading to confirm mapping.
2. Over‑augmenting data. Aggressive augmentations (e.g., extreme rotations) can produce unrealistic samples that hurt accuracy. Use tf.image.random_flip_left_right and modest rotations.
3. Ignoring class imbalance. If some classes have far fewer examples, employ class_weight in model.fit or use oversampling techniques.
4. Not normalising pixel values. Feeding raw 0‑255 values to a model expecting 0‑1 range leads to poor convergence.
5. Forgetting to clear the GPU cache. When iterating on model architecture, stale tensors can exhaust memory. Use tf.keras.backend.clear_session() between runs.
Tips and Tricks
• Use mixed precision. On modern GPUs, enable tf.keras.mixed_precision.set_global_policy('mixed_float16') to speed up training without sacrificing accuracy.
• Leverage TensorBoard. Add tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir='./logs') and pass it to model.fit for live visualisation of loss, accuracy, and learning‑rate schedules.
• Export to TensorFlow Lite for edge devices. After training, run tf.lite.TFLiteConverter.from_saved_model(export_path).convert() and deploy on Android or microcontrollers.
• Version control your data. Store raw data hashes in .gitignore and keep a data/README.md describing acquisition steps. This makes reproducibility easier for teammates.
Frequently Asked Questions
Can I train the pipeline on a CPU only?
Yes, TensorFlow runs on CPU, but training on large image datasets will be significantly slower. For prototypes or small datasets (under 5 k images) a CPU is acceptable.
How do I handle multi‑label classification?
Replace the final Dense layer with Dense(num_classes, activation='sigmoid') and use binary_crossentropy as the loss function. Adjust the dataset to yield a binary vector per sample.
What if my images have varying aspect ratios?
Resize with padding to preserve aspect ratio: use tf.image.resize_with_pad instead of plain tf.image.resize. This prevents distortion that could confuse the model.
Conclusion
Building a custom image recognition pipeline with TensorFlow involves careful data preparation, leveraging transfer learning, systematic fine‑tuning, and robust deployment via Docker and TensorFlow Serving. By following the steps above—and keeping an eye on the common pitfalls—you’ll be able to ship production‑grade models that scale from research notebooks to cloud‑native services. Remember to iterate: experiment with different backbones, augmentations, and hyper‑parameters, and always validate your changes with TensorBoard and a held‑out test set. Happy coding!
Photo by Portraits of Modern Dairy on Unsplash
