Build an ML Model with PyTorch and TensorFlow
A gentle introduction to ML and optimizing techniques
Overview
Last time, we trained a neural network by hand with pen and paper. Now it’s time to hand that work off to a software framework. Naturally, the data we’re using will be more complex, since we want to showcase what these tools can actually do.
For the XOR gate, we had 9 parameters, made up of 6 weights and 3 biases. This time, the inputs stay just as simple (still two numbers per example), but the data volume and the network both grow. We’ll use 1,000 data points instead of 4, and the network itself has 337 parameters instead of 9. Despite the larger scale, the core functionality still boils down to the same three critical pieces: input, weights, and loss function.
That basic two-layer perceptron example left us with five open questions – stochastic gradient descent, mini-batches, overfitting, dropout, early stopping. In a nutshell, they are all techniques that make training more accurate and more efficient. They’ll become far more intuitive once we’ve worked through a real example.
The Framework
Being an expert programmer isn’t a prerequisite, but some foundational Python skills are required. A good gut check is whether you can build a simple to-do list.
NOTE: There are plenty of free resources to help build that foundation first, such as this Khan Academy series.
PyTorch (from Meta) and TensorFlow (from Google) are the two frameworks used everywhere, from large-scale LLMs to ML models that run on tiny devices. Under the hood, the core workflow is the same, and it consists of the following pieces.
- Tensors – Matrices that know how to run on a GPU. A rank-2 tensor, for instance, is just a 2-dimensional matrix.
- Autograd – Automatic differentiation, so nobody has to compute \(\frac{\partial L}{\partial w}\) by hand.
- Prebuilt layers, losses, and optimizers – The pieces required to actually train a model.
One thing worth mentioning is that given a similar skill level, creativity is what differentiates an exceptional ML practitioner from an average one.
Understand the Data
The two-moons dataset looks intimidating when plotted out, but it’s just a slightly more complex version of the XOR example. XOR has simple binary inputs (1 or 0). In the case of two-moons, the inputs cover a continuous range. For example, a point at \(x_1 = 0.5\) and \(x_2 = -0.5\) would be classified as 1.
Remember how, in the XOR example, it was impossible to draw the prediction boundary with a single line (one neuron)? How many neurons and hidden layers will it take to make an accurate prediction here? By looking at the plot, you might estimate the S-curve needs 8-10 straight-line segments to trace it. However, as data gets bigger, it becomes impossible to plot and determine this by hand. Since there’s no set formula for calculating the architecture, creativity becomes relatively important. You might decide to study similar examples, then set up logic for rapid trial-and-error to home in on the optimal parameters.
Data is generated using scikit-learn library:
from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split
X, y = make_moons(n_samples=1000, noise=0.20, random_state=42)
X_train, X_val, y_train, y_val = train_test_split(
X, y, test_size=0.2, random_state=42
)Implementation
We generated 1,000 samples: 800 for the training set and 200 for the validation set. When determining accuracy, it’s important to use data the model has never seen before for the validation set.
Step 1 - Design the Network
Earlier we mentioned that there will be 337 parameters. The number comes from our architecture:
- 2 inputs times 16 neurons plus 16 biases for the first hidden layer: \(2 \times 16 + 16 = 48\)
- 16 inputs (the outputs of the first hidden layer) times 16 neurons plus 16 biases: \(16 \times 16 + 16 = 272\)
- 16 inputs times 1 neuron plus 1 bias: \(16 \times 1 + 1 = 17\)
- Total parameters: \(48 + 272 + 17 = 337\)
For example, this matrix operation will create 16 outputs for the first hidden layer: \[ \underbrace{\begin{bmatrix} x_1 & x_2 \end{bmatrix}}_{1 \times 2} \cdot \underbrace{\begin{bmatrix} w_{1,1} & w_{1,2} & \cdots & w_{1,16} \\ w_{2,1} & w_{2,2} & \cdots & w_{2,16} \end{bmatrix}}_{2 \times 16} + \underbrace{\begin{bmatrix} b_1 & b_2 & \cdots & b_{16} \end{bmatrix}}_{1 \times 16} \]
Each column of \(x\) is an input, each column of \(w\) corresponds to a neuron, and each column of \(b\) is a bias.
Step 2 - Activation Function
Previously we used the sigmoid activation function for non-linearity. It’s widely used for the output of a binary classifier. However, there’s a slight issue when we use it in the hidden layers. Say a neuron produces \(z = 2.5\), then \(\sigma(z) \approx 0.924\). During backpropagation, the local derivative at that neuron would be \(0.924\,(1 - 0.924) \approx 0.07\). Now what happens as this factor travels back through earlier layers? If each layer contributes a similar factor, the chain rule multiplies them together.
- Second layer back: \(0.07^2 = 0.0049\)
- Third layer back: \(0.07^3 = 0.000343\)
Imagine a really long backpropagation chain. The gradient reaching the earliest weights becomes vanishingly small, so those weights barely change each step according to the gradient descent formula (\(w \leftarrow w - \eta \frac{\partial L}{\partial w}\)), and the early layers take a very long time to learn anything. This is known as the vanishing gradient problem.
This is why it’s worth introducing ReLU, another activation function. It simply preserves positive values and outputs zero for negative ones, giving it a derivative of 1 for any positive input while still providing non-linearity. That means the earlier weights get a chance to update and contribute to the final prediction. It’s a good option for the hidden layers of this dataset.
\[ \mathrm{ReLU}(z) = \max(0, z) \]
The output neuron still uses sigmoid, because the prediction should be a probability between 0 and 1 for binary classification. The loss is still binary cross-entropy, the formula we previously computed by hand.
Step 3 - Putting It in Sequence
Network flow should be clear at this point:
input \(\times\) weights → ReLU → input \(\times\) weights → ReLU → input \(\times\) weights → \(\sigma\)
We already created data and separated them into training and validation set:
- Training set:
X_trainandy_train - Validation set:
X_valandy_val
Now let’s put them in sequence in PyTorch and TensorFlow side by side, because they’re quite similar.
1) Library Import
This should be pretty self explanatory. If we want to use a tool for a job, we have to make sure it’s in our toolbox.
PyTorch
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDatasetTensorFlow
import tensorflow as tf
from tensorflow.keras import layersNow it’s time to define model architecture, which include setup for weights, hidden layers, activation function, and loss function.
2) Define Model Architecture
We define each layer and it’s output by putitng them in a sequence.
PyTorch
model = nn.Sequential(
nn.Linear(2, 16), # input layer -> hidden layer 1
nn.ReLU(),
nn.Linear(16, 16), # hidden layer 1 -> hidden layer 2
nn.ReLU(),
nn.Linear(16, 1), # hidden layer 2 -> output neuron
nn.Sigmoid(), # squash to a probability, same as part one
)
loss_fn = nn.BCELoss() # binary cross-entropy
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)TensorFlow
model = tf.keras.Sequential([
layers.Input(shape=(2,)),
layers.Dense(16, activation="relu"),
layers.Dense(16, activation="relu"),
layers.Dense(1, activation="sigmoid"),
])
model.compile(
optimizer=tf.keras.optimizers.SGD(learning_rate=0.5),
loss="binary_crossentropy",
metrics=["accuracy"],
)One thing that might stand out is optimizer. .SGD in the method name stands for Stochastic Gradient Descent, whose benefit we’ll see when we get to the batch_size section, but all this does for now is just set the learning rate.
3) Use the Correct Data Structure
In PyTorch, data must be converted into tensors, the framework’s core data structure. Tensors support automatic differentiation and can run on a GPU. It’s a mandatory component for training.
PyTorch
# NumPy arrays -> float32 tensors, labels shaped as a column vector
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32).unsqueeze(1)
X_val = torch.tensor(X_val, dtype=torch.float32)
y_val = torch.tensor(y_val, dtype=torch.float32).unsqueeze(1)On the other hand TensorFlow acceps numpy or python arrays.
4) Training the Model
PyTorch:
train_loader = DataLoader(
TensorDataset(X_train, y_train), batch_size=32, shuffle=True
)
for epoch in range(100):
model.train()
for xb, yb in train_loader: # one mini-batch at a time
optimizer.zero_grad() # clear old gradients
y_hat = model(xb) # forward pass
loss = loss_fn(y_hat, yb) # how wrong are we?
loss.backward() # backpropagation
optimizer.step() # w <- w - lr * dL/dw
if (epoch + 1) % 20 == 0:
model.eval()
with torch.no_grad():
val_pred = model(X_val)
val_loss = loss_fn(val_pred, y_val)
val_acc = ((val_pred > 0.5).float() == y_val).float().mean()
print(f"epoch {epoch+1:3d} val_loss {val_loss:.4f} val_acc {val_acc:.3f}")TensorFlow:
history = model.fit(
X_train, y_train,
batch_size=32,
epochs=100,
validation_data=(X_val, y_val),
verbose=0, # silence the per-epoch progress bar
)
val_loss, val_acc = model.evaluate(X_val, y_val, verbose=0)
print(f"val_loss {val_loss:.4f} val_acc {val_acc:.3f}")One thing that should stand out is batch_size, since we haven’t discussed it. It controls how many samples are used to compute each gradient update. Given 800 samples, one epoch (one full pass through the data) can be organized three ways:
a. Pure SGD (batch_size=1) — 800 updates per epoch, 1 sample each
- Shuffle the 800 samples.
- Take sample #1 → compute its gradient → update all 337 parameters.
- Take sample #2 → compute its gradient → update all 337 parameters.
- … repeat for all 800 samples.
b. Mini-batch SGD (batch_size=32) — 25 updates per epoch, 32 samples each
- Shuffle the 800 samples, split into 25 batches of 32.
- Batch 1 (samples 1–32) → average their 32 gradients → update all 337 parameters.
- Batch 2 (samples 33–64) → average → update all 337 parameters.
- … repeat for all 25 batches.
c. Full batch (batch_size=800) — 1 update per epoch, all samples
- Take all 800 samples → average their 800 gradients → update all 337 parameters.
Pure SGD makes the most updates per epoch, but each update is noisy, because it relies on one sample, and processing samples one at a time leaves parallel hardware underused. Full batch produces the most accurate gradient per update, but only one update per epoch, so it typically needs many more epochs to converge. It also requires holding activations for the entire dataset in memory during backprop, which can exceed GPU memory and crash training on large datasets. Mini-batch strikes a balance in between. It has many update per epoch, and the batches are large enough to use hardware efficiently.
Instead of computing the exact gradient using the entire dataset, which could be millions of examples, Stochastic Gradient Descent (SDG) estimate it using small batch chosen at random to update parameters.
Verify Effectiveness
We posed a question earlier: how many lines (neurons) would we need to make a reasonable prediction? After running the code and validating accuracy, we can plot the prediction boundary to visualize the result.
That S-shaped curve shows us that two layers of ReLU neurons bend one boundary through the gap between the crescents. The few misclassified points sit where the noise genuinely mixes the classes; the error stems from messiness in the data, not something the model should compensate for.
Regularization and Optimization
Now we’re equipped with the tools and knowledge to design and run a model with less than 100 lines of code. That’s quite empowering. In practice, even given clean data, overfitting might be one of the biggest issues jeopardizing a model’s accuracy.
Intuitively, it might make sense to train as many epochs as possible to improve accuracy, but at some point, you might notice that training loss becomes stagnant and validation loss explodes. It’s often an indication that the model can only discern the training data, and has lost the ability to generalize outside of the training data.
Since this problem has been around for a long time, talented researchers have already come up with a few methods to overcome this issue.
Dropout
Dropout attacks memorization by sabotage. During training, every pass through a dropout layer switches each neuron off with probability \(p\):
PyTorch
nn.Sequential(
nn.Linear(2, 256),
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(256, 256),
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(256, 1),
nn.Sigmoid(),
)TensorFlow
tf.keras.Sequential([
layers.Dense(256, activation="relu"),
layers.Dropout(0.5),
layers.Dense(256, activation="relu"),
layers.Dropout(0.5),
layers.Dense(1, activation="sigmoid"),
])The mechanism is that each activation has about a 50% (\(p=0.5\)) chance to get dropped to zero, and the remaining activations get scaled by \(1/(1-p)\) to compensate for the dropped neurons.
This prevents a specific set of neurons from learning the majority of the task. Since activations get dropped randomly, learning can be better distributed across different neurons.
NOTE: Dropout is used during training and is inactive during inference (prediction).
Early Stopping
In a nutshell, the idea is to stop training before overfitting sets in.
PyTorch
best_val, patience, wait = float("inf"), 20, 0
for epoch in range(400):
train_one_epoch(model, train_loader)
val_loss = evaluate(model, X_val, y_val)
if val_loss < best_val:
best_val, wait = val_loss, 0
torch.save(model.state_dict(), "best.pt")
else:
wait += 1
if wait >= patience:
break
model.load_state_dict(torch.load("best.pt"))TensorFlow
early_stop = tf.keras.callbacks.EarlyStopping(
monitor="val_loss", patience=20, restore_best_weights=True
)
model.fit(X_train, y_train, epochs=400,
validation_data=(X_val, y_val), callbacks=[early_stop])Dynamic Learning Rate
We used optimizer = torch.optim.SGD(model.parameters(), lr=0.5) earlier. The issue with a fixed learning rate is that we’re stuck guessing: too small and convergence is slow, too big and the parameters hop around the minimum instead of settling. In practice, this is solved from two directions. Adaptive optimizers like Adam/AdamW maintain running statistics of each parameter’s gradients and use them to scale each parameter’s effective step size individually. A learning rate scheduler complements this by adjusting the global base rate over the course of training. For example, we might want to reduce the step size as we approach convergence.
PyTorch
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer,
factor=0.5, # halve the lr when triggered
patience=10, # wait 10 epochs of no val_loss improvement first
)