You are currently viewing Introduction to Deep Learning with TensorFlow and PyTorch in Python

Introduction to Deep Learning with TensorFlow and PyTorch in Python

  • Post author:
  • Post category:Python
  • Post comments:0 Comments
  • Post last modified:February 23, 2024

In this tutorial, we’ll introduce you to deep learning using two popular Python libraries: TensorFlow and PyTorch. We’ll cover the basics of creating neural networks, training them on data, and making predictions. By the end, you should have a good understanding of how to get started with deep learning using these powerful frameworks.


  • Basic understanding of Python
  • Familiarity with machine learning concepts is helpful but not required


Before we get started, make sure you have TensorFlow and PyTorch installed. You can install them using pip:

pip install tensorflow
pip install torch torchvision

Now, let’s dive into creating neural networks with TensorFlow and PyTorch.

1. TensorFlow

Creating a Neural Network

Let’s create a simple neural network using TensorFlow to classify handwritten digits from the MNIST dataset.

import tensorflow as tf
from tensorflow import keras

# Load the MNIST dataset
mnist = keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# Preprocess the data
train_images = train_images / 255.0
test_images = test_images / 255.0

# Build the model
model = keras.Sequential([
    keras.layers.Flatten(input_shape=(28, 28)),
    keras.layers.Dense(128, activation='relu'),
    keras.layers.Dense(10, activation='softmax')

# Compile the model

# Train the model, train_labels, epochs=5)

# Evaluate the model
test_loss, test_acc = model.evaluate(test_images, test_labels)
print(f'Test accuracy: {test_acc}')

Making Predictions

Now that we have a trained model, let’s make predictions on some test images.

predictions = model.predict(test_images)

# Print the predicted label for the first test image
print(f'Predicted label: {tf.argmax(predictions[0])}')

2. PyTorch

Creating a Neural Network

Let’s create the same neural network to classify MNIST digits using PyTorch.

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

# Define the neural network
class NeuralNet(nn.Module):
    def __init__(self):
        super(NeuralNet, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 128),
            nn.Linear(128, 10),

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

# Initialize the model
model = NeuralNet()

# Define the loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# Load the data
transform = transforms.Compose([
    transforms.Normalize((0.5,), (0.5,))

train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transform, download=True)
train_loader =, batch_size=64, shuffle=True)

# Train the model
num_epochs = 5
for epoch in range(num_epochs):
    for images, labels in train_loader:
        outputs = model(images)
        loss = criterion(outputs, labels)

# Test the model
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=transform, download=True)
test_loader =, batch_size=64, shuffle=False)

correct = 0
total = 0
with torch.no_grad():
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Test accuracy: {100 * correct / total}%')

Making Predictions

Now let’s make predictions using our trained PyTorch model.

# Get a batch of test images
dataiter = iter(test_loader)
images, labels =

# Make predictions
outputs = model(images)
_, predicted = torch.max(outputs, 1)

# Print predicted labels
print('Predicted: ', ' '.join(f'{predicted[j]}' for j in range(8)))


In this tutorial, we covered the basics of creating and training neural networks using TensorFlow and PyTorch in Python. We created simple models to classify handwritten digits from the MNIST dataset and made predictions on test images.

Deep learning is a vast field, and these examples only scratch the surface. We encourage you to explore more complex models, datasets, and techniques to further your understanding of deep learning with TensorFlow and PyTorch.

Leave a Reply