Build Your First Neural Network: A Deep Dive into the TensorFlow Library

Building your first neural network with TensorFlow is one of the most practical ways to understand how machine learning actually works under the hood. Rather than treating deep learning as a black box, walking through the process of designing, training, and evaluating a network gives you a concrete mental model you can carry into more advanced projects. TensorFlow, Google's open-source library for numerical computation and machine learning, has become one of the two dominant frameworks in the field, and its high-level Keras API makes the barrier to entry lower than most newcomers expect.

This guide walks through the full lifecycle of a simple neural network: setting up your environment, preparing data, defining the model architecture, training it, and evaluating the results. Along the way, we'll touch on the design decisions that matter and the tradeoffs you should be aware of before treating any of this as production-ready.

Why TensorFlow Is a Reasonable Starting Point

TensorFlow was released by Google in 2015 and has since matured into a comprehensive ecosystem that spans research, production deployment, and mobile or edge inference through TensorFlow Lite. Its main competitor, PyTorch, is often praised for a more Pythonic, dynamic-graph feel that some developers find easier to debug. TensorFlow's advantage lies elsewhere: its Keras API, tight integration with TensorFlow Serving and TensorFlow.js, and broad industry adoption mean that skills learned here transfer directly to deployment scenarios you're likely to encounter in a job or a shipped product.

Neither framework is objectively superior for every use case. If your work leans heavily toward research prototyping, PyTorch's flexibility often wins. If you're building something you intend to deploy at scale, in a browser, or on a mobile device, TensorFlow's tooling tends to reduce friction later in the project. It's worth knowing both exist, and this article uses TensorFlow because its Keras interface is arguably the gentlest on-ramp for someone building a first network.

Setting Up Your Environment

Before writing any model code, you need a working Python environment with TensorFlow installed. A virtual environment keeps this isolated from other projects.

python3 -m venv tf-env
source tf-env/bin/activate
pip install tensorflow

Once installed, confirm the setup works and check which version you're running, since TensorFlow's API has changed meaningfully across major versions.

import tensorflow as tf
print(tf.__version__)

If you have a compatible NVIDIA GPU, TensorFlow can use it automatically, though a CPU is perfectly sufficient for the small example we'll build here. Training a full-scale model on CPU only becomes painful once your dataset and network grow substantially larger.

Choosing a Dataset

For a first neural network, it makes sense to use a dataset that's small, clean, and well understood, so that any problems you hit are almost certainly in your code rather than in messy data. The MNIST dataset of handwritten digits is the classic choice for this reason. It contains 70,000 grayscale images of digits 0 through 9, each 28x28 pixels, split into training and test sets.

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

x_train = x_train / 255.0
x_test = x_test / 255.0

That normalization step, dividing pixel values by 255, scales them into a 0 to 1 range. Neural networks generally train faster and more reliably on normalized inputs, since large or wildly varying input magnitudes can destabilize gradient descent.

Defining the Model Architecture

With Keras, you define a network by stacking layers. For MNIST, a modest fully connected network is enough to get strong results without needing convolutional layers, though those would improve accuracy further.

model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10)
])

Each piece here does specific work. The Flatten layer converts each 28x28 image into a single 784-element vector, since dense layers expect one-dimensional input. The Dense layer with 128 units and ReLU activation is where most of the actual learning happens, combining inputs with weights and applying a nonlinearity so the network can model more than straight lines. Dropout randomly disables 20 percent of neurons during training, which is a regularization technique that helps prevent the network from memorizing the training data instead of learning generalizable patterns. The final Dense layer with 10 units corresponds to the 10 possible digit classes.

Compiling and Training the Model

Before training, the model needs an optimizer, a loss function, and a metric to track.

model.compile(
    optimizer='adam',
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy']
)

model.fit(x_train, y_train, epochs=5)

Adam is a reasonable default optimizer for most beginner projects because it adapts its learning rate automatically and tends to converge reliably without much manual tuning. Sparse categorical crossentropy is the appropriate loss function here because the labels are integers (0 through 9) rather than one-hot encoded vectors. The from_logits=True argument matters because our final layer has no activation function, meaning it outputs raw scores rather than probabilities; TensorFlow needs to know this to compute the loss correctly.

Training for 5 epochs, meaning 5 full passes through the training data, is usually enough to see this simple network reach around 97 to 98 percent accuracy on MNIST. That's not a benchmark to be impressed by in 2026, since MNIST is a solved problem for even basic architectures, but it's a useful confirmation that your pipeline actually works end to end.

Evaluating the Model

Accuracy on training data tells you almost nothing about how the model will perform on data it hasn't seen. The test set exists specifically to check this.

test_loss, test_acc = model.evaluate(x_test, y_test, verbose=2)
print('Test accuracy:', test_acc)

If test accuracy is noticeably lower than training accuracy, that gap is a signal of overfitting: the model has started memorizing quirks of the training set rather than learning generalizable patterns. On a simple MNIST model like this one, the gap should be small. On more complex projects with less data or more parameters, watching this gap closely becomes one of the more important diagnostic habits you can build.

Making Predictions

Once trained, the model can classify new images. Wrapping it with a softmax layer converts the raw logit outputs into interpretable probabilities.

probability_model = tf.keras.Sequential([
    model,
    tf.keras.layers.Softmax()
])

predictions = probability_model.predict(x_test[:5])

Each row in predictions is a probability distribution across the 10 digit classes, and the class with the highest value is the model's best guess. This distinction between logits and probabilities trips up a lot of newcomers, since it's easy to forget that raw model output and a genuine probability aren't automatically the same thing without an explicit softmax step.

Common Pitfalls Worth Knowing

A few mistakes recur often enough among people building their first network that they're worth naming directly. Forgetting to normalize input data is common and can lead to slow or unstable training. Mismatching the loss function with the label format, using categorical crossentropy on integer labels instead of sparse categorical crossentropy, throws a shape error that confuses a lot of beginners. Training for too many epochs on a small dataset without any validation monitoring can quietly produce an overfit model that looks great on paper and performs poorly on anything new.

It's also worth being honest that MNIST is not representative of most real-world problems. Real datasets are messier, class-imbalanced, and rarely arrive pre-cleaned. The habits built here- checking shapes, normalizing inputs, watching the train-test accuracy gap-matter more than the specific architecture, and they transfer directly once you move to harder problems like image classification on custom datasets or natural language tasks.

Where to Go From Here

Once this basic pipeline feels comfortable, the natural next step is experimenting with convolutional neural networks for image tasks, which handle spatial structure far better than the flatten-and-dense approach used here. From there, TensorFlow's ecosystem offers a fairly direct path toward more advanced work: TensorFlow Datasets for handling larger and more varied data sources, TensorBoard for visualizing training in more depth than print statements allow, and TensorFlow Lite or TensorFlow.js if deployment to mobile or web is eventually the goal.

Building a first neural network is less about the specific numbers you get and more about internalizing the shape of the workflow: prepare data, define architecture, compile with the right loss and optimizer, train, evaluate honestly, and iterate. That workflow doesn't change much as the problems get harder; only the pieces inside it grow more sophisticated.