Migrating from TensorFlow to PyTorch: Step-by-Step Guide

Understanding the Landscape: TensorFlow vs. PyTorch

Migrating a deep learning project from TensorFlow to PyTorch is more than a syntax swap—it’s a paradigm shift. TensorFlow (especially with Keras) provides a managed, declarative environment where models, training, and evaluation are handled by high-level APIs. PyTorch gives you explicit control over every aspect, from tensor operations to gradient updates, making it feel more like standard Python. This tutorial walks you through a complete, step-by-step migration, covering data pipelines, model architectures, training loops, checkpointing, and best practices.

What the Migration Entails

The migration involves translating:

• Data loading – from tf.data.Dataset to PyTorch DataLoader and Dataset.
• Model definition – from Keras Sequential/functional models to nn.Module subclasses.
• Training orchestration – from model.compile()/model.fit() to hand-written training loops with optimizers and loss functions.
• Checkpointing – from Keras callbacks or model.save() to saving/loading state_dicts.
• Inference – from model.evaluate()/model.predict() to manual torch.no_grad() loops.

Why Migrate?

Developers choose PyTorch for several compelling reasons:

• Dynamic computation graphs – Debugging is as simple as adding print() statements; the graph is built on-the-fly.
• Pythonic feel – PyTorch integrates naturally with Python debugging tools, numpy-like tensor operations, and standard control flow.
• Research-friendly ecosystem – Most state-of-the-art papers release PyTorch code first; frameworks like Hugging Face Transformers, PyTorch Lightning, and Detectron2 are built on it.
• Growing industry adoption – PyTorch is now the dominant framework in research and is rapidly gaining ground in production with TorchServe and ONNX export.

Core Conceptual Differences

Static vs. Dynamic Computation Graphs

TensorFlow 2.x defaults to eager execution, but many production models and older codebases still use graph mode (tf.function). PyTorch’s graph is always dynamic—you can change it at every iteration, use variable-length sequences naturally, and debug with standard Python tools. This eliminates the mental overhead of graph construction versus execution.

Keras Models vs. PyTorch Modules

In Keras, you typically build models using the Sequential API or the functional Model class. Layers are called sequentially and activations are baked into the layer call. In PyTorch, you subclass nn.Module, define layers in __init__, and implement the forward pass in forward(). Activations are applied explicitly via torch.nn.functional (e.g., F.relu()) or as separate nn.Module instances.

Step-by-Step Migration Guide

Step 1: Set Up Your PyTorch Environment

Start by installing PyTorch. Visit pytorch.org to get the exact command for your system (CUDA or CPU).

pip install torch torchvision torchaudio

Optional but recommended: install pytorch-lightning later to simplify training loops (covered in best practices).

Step 2: Convert the Data Pipeline

TensorFlow projects often rely on tf.data.Dataset for efficient input pipelines with prefetching and augmentation. PyTorch uses torch.utils.data.Dataset and DataLoader. Here’s a typical conversion:

TensorFlow data loading

import tensorflow as tf
import numpy as np

# Assume x_train, y_train are numpy arrays
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.batch(32).shuffle(1000).prefetch(tf.data.AUTOTUNE)

val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))
val_dataset = val_dataset.batch(32).prefetch(tf.data.AUTOTUNE)

PyTorch equivalent

import torch
from torch.utils.data import TensorDataset, DataLoader

# Convert numpy arrays to tensors
train_data = TensorDataset(torch.from_numpy(x_train), torch.from_numpy(y_train))
val_data = TensorDataset(torch.from_numpy(x_val), torch.from_numpy(y_val))

train_loader = DataLoader(train_data, batch_size=32, shuffle=True)
val_loader = DataLoader(val_data, batch_size=32, shuffle=False)

For complex pipelines (file reading, online augmentation), implement a custom Dataset class inheriting from torch.utils.data.Dataset, overriding __len__ and __getitem__. PyTorch’s DataLoader supports multi-processing and prefetching via num_workers and pin_memory.

Step 3: Rewrite the Model Architecture

Let’s convert a simple CNN classifier from Keras to PyTorch. Notice the key differences: channel ordering (NHWC vs. NCHW), activation placement, and the final layer’s output (Keras uses softmax; PyTorch CrossEntropyLoss expects raw logits).

TensorFlow / Keras model

import tensorflow as tf
from tensorflow.keras import layers, Model

class TFClassifier(Model):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv1 = layers.Conv2D(32, 3, padding='same', activation='relu')
        self.pool = layers.MaxPooling2D(2)
        self.flatten = layers.Flatten()
        self.fc1 = layers.Dense(128, activation='relu')
        self.fc2 = layers.Dense(num_classes, activation='softmax')

    def call(self, inputs):
        x = self.conv1(inputs)
        x = self.pool(x)
        x = self.flatten(x)
        x = self.fc1(x)
        return self.fc2(x)

model = TFClassifier()

PyTorch equivalent

import torch
import torch.nn as nn
import torch.nn.functional as F

class PyTorchClassifier(nn.Module):
    def __init__(self, input_channels=3, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2d(input_channels, 32, 3, padding=1)  # 'same' padding
        self.pool = nn.MaxPool2d(2)
        # Calculate flattened size after conv+pool (example for 32x32 input)
        self.fc1 = nn.Linear(32 * 16 * 16, 128)   # Adjust based on actual input size
        self.fc2 = nn.Linear(128, num_classes)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))  # Apply conv then relu, then pool
        x = x.view(x.size(0), -1)             # Flatten
        x = F.relu(self.fc1(x))
        x = self.fc2(x)                       # Raw logits, no softmax
        return x

model = PyTorchClassifier()

Key takeaways:

• PyTorch layers expect input shape (batch, channels, height, width). If your TensorFlow data is in NHWC format, use torch.permute() or reshape before feeding.
• Activations like ReLU are applied explicitly; do not include them inside nn.Conv2d or nn.Linear.
• The final layer should output raw logits when using nn.CrossEntropyLoss.

Step 4: Adapt the Training Loop

Keras abstracts training with model.compile() and model.fit(). PyTorch requires you to write the loop yourself—this gives you complete flexibility but demands careful handling of gradients and device placement.

TensorFlow training

model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)
model.fit(train_dataset, epochs=5, validation_data=val_dataset)

PyTorch training

import torch.optim as optim

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

optimizer = optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.CrossEntropyLoss()  # Combines log_softmax + NLLLoss, expects raw logits

num_epochs = 5

for epoch in range(num_epochs):
    # Training phase
    model.train()
    running_loss = 0.0
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()               # Clear gradients
        output = model(data)                # Forward pass
        loss = loss_fn(output, target)      # Compute loss
        loss.backward()                     # Backpropagation
        optimizer.step()                    # Update weights
        running_loss += loss.item()

    # Validation phase
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():                   # Disable gradient computation
        for data, target in val_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            _, predicted = torch.max(output, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()

    epoch_loss = running_loss / len(train_loader)
    epoch_acc = correct / total
    print(f'Epoch {epoch+1}/{num_epochs} | Loss: {epoch_loss:.4f} | Val Acc: {epoch_acc:.4f}')

Notice the explicit device transfer (.to(device)), gradient zeroing, and the torch.no_grad() context during validation.

Step 5: Handling Checkpoints and Saving/Loading

TensorFlow provides model.save() for full model persistence. PyTorch saves only the learned parameters (state_dict) by default, which is lighter and more flexible.

TensorFlow

# Save entire model
model.save('my_model.h5')
# Load
restored_model = tf.keras.models.load_model('my_model.h5')

PyTorch

# Save only model weights
torch.save(model.state_dict(), 'model_weights.pth')

# Load weights into the same architecture
model.load_state_dict(torch.load('model_weights.pth'))

# Save a full checkpoint including optimizer state (for resuming)
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': running_loss,
}
torch.save(checkpoint, 'checkpoint.pth')

# Resume training
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch'] + 1

Step 6: Evaluation and Inference

Keras offers model.evaluate() and model.predict(). In PyTorch, you manually iterate the test DataLoader inside a torch.no_grad() block.

TensorFlow evaluation

test_loss, test_acc = model.evaluate(test_dataset)
predictions = model.predict(test_dataset)

PyTorch evaluation and inference

model.eval()
test_loss = 0.0
correct = 0
total = 0

with torch.no_grad():
    for data, target in test_loader:
        data, target = data.to(device), target.to(device)
        output = model(data)
        loss = loss_fn(output, target)
        test_loss += loss.item()
        _, predicted = torch.max(output, 1)
        total += target.size(0)
        correct += (predicted == target).sum().item()

avg_loss = test_loss / len(test_loader)
accuracy = correct / total
print(f'Test Loss: {avg_loss:.4f} | Accuracy: {accuracy:.4f}')

# For inference on new data
with torch.no_grad():
    sample_batch = torch.tensor(some_new_data).to(device)
    predictions = model(sample_batch)
    predicted_classes = torch.argmax(predictions, dim=1)

Step 7: Transfer Learning and Pre-trained Models

Both frameworks offer easy access to pre-trained models. In TensorFlow, you use tf.keras.applications. In PyTorch, torchvision.models provides architectures like ResNet, VGG, etc. Hugging Face Transformers is also widely used.

TensorFlow transfer learning

base_model = tf.keras.applications.ResNet50(include_top=False, weights='imagenet')
x = layers.GlobalAveragePooling2D()(base_model.output)
output = layers.Dense(num_classes, activation='softmax')(x)
model = Model(inputs=base_model.input, outputs=output)

PyTorch transfer learning

import torchvision.models as models

# Load pre-trained ResNet50 without the original classifier
resnet = models.resnet50(pretrained=True)
num_ftrs = resnet.fc.in_features
resnet.fc = nn.Linear(num_ftrs, num_classes)  # Replace final layer
model = resnet.to(device)

When using PyTorch, remember to freeze or set different learning rates for pre-trained layers via param.requires_grad = False and optimizer parameter groups.

Common Pitfalls and Best Practices

Batch Dimension Handling

The most common migration pitfall is channel ordering. TensorFlow often uses NHWC (batch, height, width, channels), while PyTorch uses NCHW (batch, channels, height, width). Convert data either when creating tensors or via .permute(0, 3, 1, 2). If you use torchvision.transforms.ToTensor(), it automatically converts PIL images to NCHW.

Device Management

In TensorFlow, device placement is largely automatic. PyTorch requires explicit moves to the GPU. Always define a device variable and call model.to(device) and tensor.to(device). Use pin_memory=True in DataLoader for faster CPU→GPU transfers.

Gradient Clipping and Custom Layers

PyTorch provides torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm) for gradient clipping. Custom layers are as simple as subclassing nn.Module and implementing forward(). You can mix dynamic Python constructs (loops, conditionals) directly in the forward pass.

Using PyTorch Lightning for Cleaner Code

If you miss the Keras-like abstraction for training loops, PyTorch Lightning is the recommended bridge. It encapsulates the training loop, checkpointing, logging, and multi-GPU support while keeping full PyTorch flexibility.

import pytorch_lightning as pl

class LightningClassifier(pl.LightningModule):
    def __init__(self, model, lr=0.001):
        super().__init__()
        self.model = model
        self.lr = lr
        self.loss_fn = nn.CrossEntropyLoss()

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self.model(x)
        loss = self.loss_fn(logits, y)
        self.log('train_loss', loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self.model(x)
        loss = self.loss_fn(logits, y)
        preds = torch.argmax(logits, dim=1)
        acc = (preds == y).float().mean()
        self.log('val_acc', acc, prog_bar=True)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.lr)

# Usage
model = PyTorchClassifier()
lightning_module = LightningClassifier(model)
trainer = pl.Trainer(max_epochs=5, accelerator='auto')
trainer.fit(lightning_module, train_loader, val_loader)

Conclusion

Migrating from TensorFlow to PyTorch is a strategic move toward greater flexibility, transparency, and alignment with modern research. By systematically converting your data pipeline, model definition, training loop, and checkpointing, you retain complete control while unlocking PyTorch’s dynamic debugging and extensive ecosystem. Embrace the explicit training loop—it may feel verbose at first, but it gives you the power to implement any custom logic seamlessly. Follow the best practices around device management and channel ordering, and consider PyTorch Lightning for a cleaner, scalable codebase. The transition is an investment that pays off in maintainable, production-ready deep learning projects.