How to Build Neural Networks with TensorFlow

TensorFlow is Google's open-source machine learning framework that powers everything from image recognition to natural language processing. Understanding neural networks has become essential for modern developers as AI capabilities become integrated into more applications. This guide covers building practical neural networks with TensorFlow from a senior developer's perspective.

Why TensorFlow

TensorFlow excels because:

1. Production-ready: Scales from laptop to data center
2. Ecosystem: Keras, TensorBoard, TFLite, TF.js
3. Hardware Support: CPU, GPU, TPU acceleration
4. Deployment Options: Server, mobile, browser, edge devices
5. Community: Extensive documentation and examples

Setup

Installation

# Basic installation
pip install tensorflow
# With GPU support (requires CUDA)
pip install tensorflow[and-cuda]
# For Apple Silicon Macs
pip install tensorflow-macos tensorflow-metal
# Additional useful packages
pip install pandas matplotlib scikit-learn

Verify Installation

import tensorflow as tf
print(f"TensorFlow version: {tf.__version__}")
print(f"GPU available: {tf.config.list_physical_devices('GPU')}")

GPU Setup (Optional)

For NVIDIA GPUs:

# Install the CUDA toolkit (Ubuntu)
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb
sudo dpkg -i cuda-keyring_1.0-1_all.deb
sudo apt update
sudo apt install cuda-toolkit-12-0
# Verify
nvidia-smi

Neural Network Basics

Core Concepts

• Neurons: Basic units that apply weights, bias, and activation
• Layers: Collections of neurons
• Activation Functions: Non-linear transformations (ReLU, sigmoid)
• Loss Function: Measures prediction error
• Optimizer: Updates weights to minimize loss
• Epochs: Complete passes through training data

Simple Model Structure

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
# Sequential model: layers stacked linearly
model = keras.Sequential([
layers.Dense(128, activation='relu', input_shape=(784,)),
layers.Dropout(0.2),
layers.Dense(64, activation='relu'),
layers.Dense(10, activation='softmax')
])
# Compile: configure training
model.compile(
optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy']
)
# Summary
model.summary()

Image Classification (MNIST)

Load and Prepare Data

import tensorflow as tf
from tensorflow import keras
import numpy as np
import matplotlib.pyplot as plt
# Load MNIST dataset
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
# Normalize pixel values to 0-1
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0
# Reshape for dense layers: (28, 28) -> (784,)
x_train_flat = x_train.reshape(-1, 784)
x_test_flat = x_test.reshape(-1, 784)
print(f"Training samples: {len(x_train)}")
print(f"Test samples: {len(x_test)}")
print(f"Image shape: {x_train[0].shape}")

Build and Train Model

# Create model
model = keras.Sequential([
layers.Dense(256, activation='relu', input_shape=(784,)),
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']
)
# Train
history = model.fit(
x_train_flat, y_train,
epochs=10,
batch_size=128,
validation_split=0.1,
verbose=1
)
# Evaluate
test_loss, test_accuracy = model.evaluate(x_test_flat, y_test)
print(f"Test accuracy: {test_accuracy:.4f}")

Plot Training History

plt.figure(figsize=(12, 4))
# Accuracy
plt.subplot(1, 2, 1)
plt.plot(history.history['accuracy'], label='Training')
plt.plot(history.history['val_accuracy'], label='Validation')
plt.title('Model Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
# Loss
plt.subplot(1, 2, 2)
plt.plot(history.history['loss'], label='Training')
plt.plot(history.history['val_loss'], label='Validation')
plt.title('Model Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.tight_layout()
plt.savefig('training_history.png')
plt.show()

Convolutional Neural Networks (CNN)

For image data, CNNs outperform dense networks:

# Reshape for CNN: add channel dimension
x_train_cnn = x_train.reshape(-1, 28, 28, 1)
x_test_cnn = x_test.reshape(-1, 28, 28, 1)
# Build CNN
cnn_model = keras.Sequential([
# Convolutional layers
layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
layers.MaxPooling2D((2, 2)),
layers.Conv2D(64, (3, 3), activation='relu'),
layers.MaxPooling2D((2, 2)),
layers.Conv2D(64, (3, 3), activation='relu'),
# Dense layers
layers.Flatten(),
layers.Dense(64, activation='relu'),
layers.Dropout(0.5),
layers.Dense(10, activation='softmax')
])
cnn_model.compile(
optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy']
)
cnn_model.fit(
x_train_cnn, y_train,
epochs=5,
batch_size=64,
validation_split=0.1
)
# CNNs typically achieve 99%+ accuracy on MNIST
cnn_model.evaluate(x_test_cnn, y_test)

Transfer Learning

Use pre-trained models for better results with less data:

from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.preprocessing.image import ImageDataGenerator
# Load pre-trained model (without top classification layer)
base_model = MobileNetV2(
weights='imagenet',
include_top=False,
input_shape=(224, 224, 3)
)
# Freeze base model weights
base_model.trainable = False
# Add custom classification head
model = keras.Sequential([
base_model,
layers.GlobalAveragePooling2D(),
layers.Dense(256, activation='relu'),
layers.Dropout(0.5),
layers.Dense(10, activation='softmax') # 10 classes
])
model.compile(
optimizer=keras.optimizers.Adam(learning_rate=0.001),
loss='categorical_crossentropy',
metrics=['accuracy']
)
# Data augmentation for training
train_datagen = ImageDataGenerator(
rescale=1./255,
rotation_range=20,
width_shift_range=0.2,
height_shift_range=0.2,
horizontal_flip=True,
validation_split=0.2
)
# Load your own images
train_generator = train_datagen.flow_from_directory(
'data/train',
target_size=(224, 224),
batch_size=32,
class_mode='categorical',
subset='training'
)
validation_generator = train_datagen.flow_from_directory(
'data/train',
target_size=(224, 224),
batch_size=32,
class_mode='categorical',
subset='validation'
)
# Train
model.fit(
train_generator,
epochs=10,
validation_data=validation_generator
)

Saving and Loading Models

# Save entire model
model.save('my_model.keras')
# Load model
loaded_model = keras.models.load_model('my_model.keras')
# Save weights only
model.save_weights('my_model_weights.weights.h5')
# Load weights
model.load_weights('my_model_weights.weights.h5')
# Export for TensorFlow Serving
model.export('saved_model_dir')
# Convert to TensorFlow Lite (mobile/edge)
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
with open('model.tflite', 'wb') as f:
f.write(tflite_model)

Making Predictions

# Single prediction
image = x_test[0:1] # Keep batch dimension
prediction = model.predict(image)
predicted_class = np.argmax(prediction[0])
confidence = prediction[0][predicted_class]
print(f"Predicted: {predicted_class}, Confidence: {confidence:.2%}")
# Batch prediction
predictions = model.predict(x_test[:100])
predicted_classes = np.argmax(predictions, axis=1)
# Evaluate predictions
from sklearn.metrics import classification_report, confusion_matrix
print(classification_report(y_test[:100], predicted_classes))

Custom Training Loop

For more control:

# Custom training
model = keras.Sequential([
layers.Dense(128, activation='relu'),
layers.Dense(10, activation='softmax')
])
loss_fn = keras.losses.SparseCategoricalCrossentropy()
optimizer = keras.optimizers.Adam()
# Training metrics
train_loss = keras.metrics.Mean()
train_accuracy = keras.metrics.SparseCategoricalAccuracy()
@tf.function
def train_step(x, y):
with tf.GradientTape() as tape:
predictions = model(x, training=True)
loss = loss_fn(y, predictions)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
train_loss.update_state(loss)
train_accuracy.update_state(y, predictions)
# Training loop
epochs = 10
batch_size = 32
dataset = tf.data.Dataset.from_tensor_slices((x_train_flat, y_train))
dataset = dataset.shuffle(10000).batch(batch_size)
for epoch in range(epochs):
train_loss.reset_states()
train_accuracy.reset_states()
for x_batch, y_batch in dataset:
train_step(x_batch, y_batch)
print(f"Epoch {epoch + 1}: Loss = {train_loss.result():.4f}, "
f"Accuracy = {train_accuracy.result():.4f}")

TensorBoard for Visualization

from tensorflow.keras.callbacks import TensorBoard
import datetime
# Create log directory
log_dir = "logs/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = TensorBoard(
log_dir=log_dir,
histogram_freq=1,
write_graph=True,
write_images=True
)
# Train with callback
model.fit(
x_train_flat, y_train,
epochs=10,
validation_split=0.1,
callbacks=[tensorboard_callback]
)
# Launch TensorBoard
# tensorboard --logdir logs/

Practical Tips

Prevent Overfitting

# Early stopping
early_stop = keras.callbacks.EarlyStopping(
monitor='val_loss',
patience=3,
restore_best_weights=True
)
# Learning rate reduction
reduce_lr = keras.callbacks.ReduceLROnPlateau(
monitor='val_loss',
factor=0.2,
patience=2,
min_lr=0.0001
)
model.fit(
x_train, y_train,
epochs=100,
validation_split=0.2,
callbacks=[early_stop, reduce_lr]
)

Model Checkpoint

checkpoint = keras.callbacks.ModelCheckpoint(
'best_model.keras',
monitor='val_accuracy',
save_best_only=True,
mode='max'
)

Key Takeaways

1. Start simple: Dense networks before CNNs
2. Normalize inputs: Scale to 0-1 or standardize
3. Use validation split: Monitor for overfitting
4. Transfer learning: Pre-trained models for image tasks
5. GPU acceleration: Essential for larger models
6. Save checkpoints: Don't lose training progress

Neural networks are powerful but require experimentation—start with working examples and iterate from there.