PyTorch vs TensorFlow: which deep learning framework should you use?

import torch
import torch.nn as nn
from torchvision import models
import os

# Load pre-trained model (eager execution: no graph building)
model = models.resnet50(pretrained=True)
model.train()  # Set to training mode

# Dummy batch of images (batch_size=4, 3 channels, 224x224)
x = torch.randn(4, 3, 224, 224)
y = torch.randint(0, 1000, (4,))

# Forward pass: executed immediately in eager mode
logits = model(x)
loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(logits, y)

# Backward pass and optimization
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
optimizer.zero_grad()
loss.backward()  # Compute gradients: intuitive for debugging
optimizer.step()

print(f"Loss: {loss.item():.4f}")  # Easy to inspect scalar values

# Inference (no graph conversion needed)
model.eval()
with torch.no_grad():
    predictions = model(x)  # Same code, different mode
    print(f"Predictions shape: {predictions.shape}")

import tensorflow as tf
from tensorflow.keras import applications, losses, optimizers
import os

# Load pre-trained model
model = applications.ResNet50(weights='imagenet')

# Dummy batch of images (batch_size=4, 224x224, 3 channels)
x = tf.random.normal((4, 224, 224, 3))
y = tf.random.uniform((4,), minval=0, maxval=1000, dtype=tf.int32)

# Define loss and optimizer
loss_fn = losses.SparseCategoricalCrossentropy()
optimizer = optimizers.Adam(learning_rate=1e-4)

# Training step wrapped in GradientTape (TF's graph-building abstraction)
with tf.GradientTape() as tape:
    logits = model(x, training=True)  # training flag controls BatchNorm/Dropout behavior
    loss = loss_fn(y, logits)

# Compute gradients
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))

print(f"Loss: {loss.numpy():.4f}")  # Must convert to NumPy to inspect

# Inference (same model, different call signature)
predictions = model(x, training=False)  # Explicit training=False required
print(f"Predictions shape: {predictions.shape}")