Neural Networks with PyTorch and Lightning AI Part 5: Final Results and GPU Acceleration

#machinelearning #ai

In the previous article, we saw how we automated several manual pieces when training a neural network. In this article, we will check the final results and also see how we can do hardware acceleration.

For verifying that the final_bias is correctly optimized, we can print out the new value:

trainer.fit(model, train_dataloaders=dataloader)
print(model.final_bias.data)

This gives us -16.0098.

Plotting the Result

Let’s also visualize it:

output_values = model(input_doses)

sns.set(style="whitegrid")

sns.lineplot(
    x=input_doses,
    y=output_values.detach(),
    color="green",
    linewidth=2.5
)

plt.ylabel("Effectiveness")
plt.xlabel("Dose")

A Practical Question: Why Not Always Use GPUs?

Now there are some practical questions, like: if we want to use GPU acceleration, how is that possible?

On most basic laptops, the CPU does all the work.

In that case, both:

the neural network (weights and biases)
the training data tensors

are stored and processed on the CPU.

All computations happen there.

But CPUs have limited cores, so when we run something like neural network training, it can become very slow.

Why GPUs Are Used

To speed things up, we use Graphics Processing Units (GPUs).

A GPU has 10x to 100x more compute cores compared to a CPU.

However, without tools like Lightning, using GPUs requires manually assigning tensors and operations to the correct device, which becomes complicated.

How Lightning Helps

With Lightning, this process is automated.

We can simply write:

trainer = L.Trainer(
    max_epochs=34,
    accelerator="auto",
    devices="auto"
)

Here:

accelerator="auto" allows Lightning to automatically detect whether a GPU is available
devices="auto" lets Lightning decide how many GPUs to use

Wrapping up

This means the same code can:

run on a CPU without changes
automatically scale to GPUs when available
avoid manual device management entirely

So Lightning makes training not only simpler, but also hardware-flexible without extra effort.

In the coming series, we will explore more such applications of lightning.

AI agents write code fast. They also silently remove logic, change behavior, and introduce bugs -- without telling you. You often find out in production.

git-lrc fixes this. It hooks into git commit and reviews every diff before it lands. 60-second setup. Completely free.

Any feedback or contributors are welcome! It's online, source-available, and ready for anyone to use.

Give it a ⭐ star on Github