DEV Community

Cover image for Building Your First AI Model with Python and TensorFlow
Gulshan Yadav
Gulshan Yadav

Posted on Originally published at misar.blog

Building Your First AI Model with Python and TensorFlow

A practical, honest path from zero to your first working model — the Keras Sequential API, a real data pipeline, and the production lessons that tutorials skip.

I spent a week once helping a team of interns build their first real model. They had done the coursework — linear regression in a notebook, a copied MNIST example, all the theory. On day one, I asked them to load their own dataset and get a model to train. By day three, every single one of them was stuck on the same three problems, and none of the problems was the neural network itself. The data was not normalized. The shapes did not line up. And the model they were training on 60,000 images was downloading and processing them one by one, a fresh read from disk for every batch, which made a 10-minute job take three hours.

That week taught me something I now tell every beginner: building a model is the easy part. Building a model that actually trains, evaluates honestly, and survives contact with real data is the actual skill. This guide is that skill — the exact path I use to take someone from a blank notebook to a working, trustworthy first model with Python and TensorFlow, including the failures you will hit along the way.

The framing: what you are actually building

Before any code, the mental model. A neural network is a function with a very large number of adjustable parameters. Training is the process of adjusting those parameters so the function maps inputs to correct outputs. TensorFlow is the engine that runs this process efficiently — it builds the computation graph, computes gradients, and updates the parameters. Keras, which ships inside TensorFlow, is the high-level API that lets you describe the network in plain Python instead of tensor math.

The Sequential API is where every beginner should start. You list the layers in order, top to bottom, and Keras chains them together. It cannot express exotic architectures — you will graduate to the functional API and tf.keras subclasses when you need branching or shared layers — but for a first model, Sequential is correct, and it keeps every moving part visible.

The data pipeline: the part tutorials rush

Every beginner tutorial skips straight to the model. That is a mistake, because in my experience 70% of first-model failures come from the data, not the network. Here is the pipeline you actually need, using the canonical Fashion MNIST dataset — clothing images in ten classes, which is a better first dataset than handwritten digits because it is harder, and harder teaches more.

import tensorflow as tf

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0

# Add the channel dimension: (60000, 28, 28) -> (60000, 28, 28, 1)
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]

dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
dataset = dataset.shuffle(10_000).batch(128).prefetch(tf.data.AUTOTUNE)
Enter fullscreen mode Exit fullscreen mode

Three details here are the difference between smooth and painful:

  1. Normalize to [0, 1] by dividing by 255. Pixel values from 0 to 255 are a nightmare for gradient-based training — large inputs produce large gradients and unstable updates. Scaling to floats in [0, 1] is not optional polish; it is the difference between a model that trains and one that stalls.
  2. The channel dimension. Convolutional layers expect (height, width, channels). The raw data comes in as (28, 28), so the [..., tf.newaxis] trick adds the single grayscale channel. Missing this produces the most common shape error in all of TensorFlow — ValueError: Input 0 of layer "conv2d" is incompatible. When you see that, this line is why.
  3. tf.data with shuffle, batch, and prefetch. This is the fix for the interns' three-hour training run. shuffle(10_000) randomizes the order so the model does not learn the order of the data, batch(128) groups samples for efficient matrix math, and prefetch(AUTOTUNE) overlaps data loading with training so the GPU never waits on disk. This pipeline is production-shaped from day one.

The model: a real CNN, not a toy

Now the network. I start beginners on a small convolutional model, not a dense MLP, because convolution is the correct tool for images — it is what the task actually requires, and learning it first saves you the "why is my accuracy stuck at 91%?" conversation later.

from tensorflow import keras
from tensorflow.keras import layers

model = keras.Sequential([
    layers.Input(shape=(28, 28, 1)),
    layers.Conv2D(32, 3, activation="relu"),
    layers.MaxPooling2D(2),
    layers.Conv2D(64, 3, activation="relu"),
    layers.MaxPooling2D(2),
    layers.Flatten(),
    layers.Dropout(0.3),
    layers.Dense(128, activation="relu"),
    layers.Dropout(0.3),
    layers.Dense(10, activation="softmax"),
])

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

Let me walk through the choices, because every line is a decision:

  • Input(shape=(28, 28, 1)) declares the expected shape explicitly. Declaring it means Keras can check your data against it and give you a clear error instead of a cryptic one.
  • Two conv blocks of 32 then 64 filters, each followed by max pooling. The pattern — more filters, smaller spatial size, deeper in the network — is the classic CNN design I covered in my architectures guide. It is not arbitrary.
  • Flatten collapses the 2D feature maps into a 1D vector so the dense layers can consume them.
  • Dropout(0.3) before the dense layers. This is your first line of defense against overfitting — during training, 30% of the neurons are randomly dropped per step, which forces the network to learn redundant, robust features instead of memorizing the training set.
  • softmax on the final dense layer turns the ten raw scores into a probability distribution across the ten classes. The predicted class is the index of the highest probability.
  • sparse_categorical_crossentropy is the loss for integer labels (0–9) with one class per sample. If your labels are one-hot encoded vectors, you would use categorical_crossentropy instead — a classic beginner mix-up.

Training and evaluating honestly

Here is where beginners go wrong in a specific way: they report training accuracy as if it meant something. Training accuracy is the model grading its own homework. The evaluation that matters is on the test set — data the model has never seen.

history = model.fit(
    dataset,                 # the tf.data pipeline
    validation_data=(x_test, y_test),
    epochs=15,
)

test_loss, test_acc = model.evaluate(x_test, y_test, verbose=0)
print(f"Test accuracy: {test_acc:.4f}")   # expect ~0.92 with this architecture
Enter fullscreen mode Exit fullscreen mode

A well-run version of this model lands around 92% test accuracy on Fashion MNIST. If you see numbers near that, the pipeline is working. If you see 99%, your test data leaked into training somewhere. If you see 70%, something is broken — and the fixes are almost never in the model.

The validation_data argument is the honest part of the loop: every epoch, the model is evaluated on the test set, and history records both training and validation metrics. Plot them side by side. The training curve rising while validation flattens or falls is the overfitting signature — and you now have a model, and you know exactly how to see when it is memorizing instead of learning.

The training loop, explained once

Beginners treat model.fit() as magic. It is not — and understanding it protects you. What actually happens, per epoch:

  1. Keras shuffles and iterates over your tf.data pipeline in batches of 128.
  2. For each batch, the model runs a forward pass: pixels in, predictions out.
  3. The loss function compares predictions to the true labels.
  4. TensorFlow computes gradients of the loss with respect to every parameter, via automatic differentiation on the computation graph.
  5. The optimizer — Adam here — takes a step: it adjusts every parameter in the direction that reduces the loss, using the gradients.
  6. Repeat until every batch is consumed. That is one epoch. Then again.

Adam is the default because it is nearly impossible to misconfigure — it adapts the learning rate per parameter, which is why it "just works" where plain SGD needs careful tuning. Fifteen epochs here is about the right amount; watch the validation curve and stop when it stops improving, because every epoch past that point is wasted compute and growing overfit.

Production reality: what happens after the notebook

The model works in a notebook. Here is what happens next, from deployments that actually shipped:

  1. Saving and reloading are not trivial. You must save the whole model, not just the weights, so the architecture comes back intact: model.save("fashion_cnn.keras") and later loaded = keras.models.load_model("fashion_cnn.keras"). Saving only model.get_weights() and rebuilding by hand is how you lose a week to a mismatched architecture.
  2. Inference on new data needs the same preprocessing. The model expects floats in [0, 1] with a channel dimension, because that is what it trained on. Feeding raw 0–255 ints into the saved model silently degrades predictions. Wrap the preprocessing in the same function you trained with, or better, bake it into a tf.keras.Sequential with a tf.keras.layers.Rescaling(1./255) layer as the first layer, so the saved model is self-contained.
  3. The metric that matters changes at deployment. In a notebook, accuracy is the number. In production, the cost of a wrong answer changes everything: a medical triage model and a clothing classifier have very different tolerances for false positives. Before you ship, decide what a mistake actually costs, and optimize for that — which usually means tracking precision and recall per class, not the headline accuracy.
  4. Retraining on your own data is where beginners actually work. The Fashion MNIST pipeline is the exercise. Your first real project is the same pipeline against your own data — a CSV, a folder of images, an API dump — and that is where you will use tf.data.Dataset.from_tensor_slices and, when the data gets large, tf.keras.utils.image_dataset_from_directory and tf.data.experimental.make_csv_dataset. The skills transfer. The data does not.

The failure modes I have seen (so you do not have to)

A short list of the real failures, in order of how often they appear:

  1. Shape errors. The channel-dimension mistake above accounts for a huge share of first-model crashes. When you see a shape error, trace it with print(x.shape) at every stage instead of guessing.
  2. Training accuracy high, test accuracy low. That is overfitting. Add dropout, add more data, or make the model smaller — in that order.
  3. Loss stuck flat from epoch one. Usually a preprocessing bug — data not normalized, or labels mismatched with the loss function. Check that inputs are in [0, 1] and the loss matches your label format.
  4. The model learns nothing and you cannot tell why. The failure is often the data, not the network: too few samples, label errors, or a pipeline that silently drops rows. Spend your debugging time on the data first.
  5. Infinite training. You are training on 60,000 images one at a time, and the machine is thrashing. That is what tf.data and prefetch fix.

When NOT to use TensorFlow for this

The uncomfortable truth: a first model is not the right tool for every first problem.

  • If your task is tabular data and interpretability matters — a churn model, a pricing model, a fraud score — a gradient-boosted tree (LightGBM or XGBoost) will usually beat a neural network on accuracy and give you feature importance you can explain to a stakeholder. I ship more of those than neural networks, and I am not ashamed of it.
  • If your data is text — start with a pretrained transformer from transformers on top of PyTorch rather than training embeddings from scratch. Training your own NLP model from zero when a pretrained one exists is burning compute you cannot afford.
  • If your dataset is a few thousand rows — a neural network will overfit. A shallow model with good features, or a tree ensemble, is the honest engineering choice.

TensorFlow with Keras is the right tool when you have real data volume, a grid-structured problem like images or sequences, and you want a system that scales to production. Match the tool to the problem, not the other way around.

The first-model checklist

Run through this before you call your first model done:

  • [ ] Data normalized to [0, 1] and cast to float32
  • [ ] Channel dimension present for image inputs
  • [ ] Data shuffled, batched, and prefetched via tf.data
  • [ ] Input(shape=...) declared so shape errors are clear
  • [ ] Dropout present if the model is dense-heavy
  • [ ] Loss matches the label format (sparse_categorical_crossentropy for integers, categorical_crossentropy for one-hot)
  • [ ] Validation accuracy reported, not just training accuracy
  • [ ] Training and validation curves plotted side by side
  • [ ] Full model saved with model.save(), not just weights
  • [ ] Preprocessing baked into the saved model so inference is self-contained
  • [ ] The deployment metric chosen (precision, recall, or accuracy) matches what a mistake actually costs

A closing thought from the interns' week

The interns finished. By day five, each of them had a working model on their own data — a normalizer pipeline, a CNN, an honest test evaluation, a saved artifact. None of what they built was exotic. It was all the pipeline in this guide: data preparation, a Sequential model, a training loop, honest evaluation. That is the entire secret of building your first AI model. Not genius, not a better architecture than anyone else's. Just the boring, correct pipeline, executed without skipping the parts that do not produce a satisfying chart.

Build this exact model tonight. It takes less than an hour on a laptop. Then change one thing — swap the dataset, add a layer, remove the dropout — and watch what happens. The model will train, the curves will tell you the truth, and by the time you have done it twice, you will not be a beginner anymore. You will be an engineer with a repeatable system.


*Gulshan Yad

Top comments (0)