PyTorch

graph TD
    A[PyTorch] --> B[Core Components]
    A --> C[Libraries & Extensions]
    A --> D[Hardware Acceleration]
    A --> E[Deployment]

    B --> B1[Tensor]
    B --> B2[Autograd]
    B --> B3[NN Module]
    B --> B4[Optimizers]

    C --> C1[TorchVision]
    C --> C2[TorchText]
    C --> C3[TorchAudio]
    C --> C4[TorchScript]
    C --> C5[JIT Compiler]

    D --> D1[CPU]
    D --> D2[GPU]
    D --> D3[TPU]
    D --> D4[Mobile/Edge]

    E --> E1[TorchScript]
    E --> E2[TorchServe]
    E --> E3[ONNX Export]
    E --> E4[Mobile Deployment]

    style A fill:#6b72ff,stroke:#333
    style B fill:#4ecdc4,stroke:#333
    style C fill:#f9ca24,stroke:#333
    style D fill:#ff6b6b,stroke:#333
    style E fill:#6c5ce7,stroke:#333

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

# 1. Set device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

# 2. Load and prepare data
print("Loading and preparing data...")
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

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

test_set = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=64, shuffle=False)

# 3. Define the model
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout = nn.Dropout(0.25)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = nn.functional.relu(x)
        x = self.conv2(x)
        x = nn.functional.relu(x)
        x = nn.functional.max_pool2d(x, 2)
        x = self.dropout(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = nn.functional.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x

model = Net().to(device)

# 4. Define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())

# 5. Train the model
print("Training the model...")
train_losses = []
train_accuracies = []

for epoch in range(5):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        _, predicted = output.max(1)
        total += target.size(0)
        correct += predicted.eq(target).sum().item()

        if batch_idx % 100 == 99:
            print(f'Epoch: {epoch+1}, Batch: {batch_idx+1}, Loss: {running_loss/100:.3f}')
            running_loss = 0.0

    epoch_loss = running_loss / len(train_loader)
    epoch_acc = 100. * correct / total
    train_losses.append(epoch_loss)
    train_accuracies.append(epoch_acc)

    print(f'Epoch {epoch+1}: Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc:.2f}%')

# 6. Evaluate the model
print("Evaluating the model...")
model.eval()
test_loss = 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)
        test_loss += criterion(output, target).item()
        _, predicted = output.max(1)
        total += target.size(0)
        correct += predicted.eq(target).sum().item()

test_loss /= len(test_loader)
test_acc = 100. * correct / total

print(f'Test Loss: {test_loss:.4f}, Test Accuracy: {test_acc:.2f}%')

# 7. Make predictions
print("Making predictions...")
dataiter = iter(test_loader)
images, labels = next(dataiter)
images, labels = images.to(device), labels.to(device)

outputs = model(images[:5])
_, predicted = torch.max(outputs, 1)

print("Sample predictions:")
for i in range(5):
    print(f'Predicted: {predicted[i].item()}, True: {labels[i].item()}')

# 8. Visualize training history
plt.figure(figsize=(12, 4))

plt.subplot(1, 2, 1)
plt.plot(train_accuracies, label='Training Accuracy')
plt.title('Training Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy (%)')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(train_losses, label='Training Loss')
plt.title('Training Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()

plt.tight_layout()
plt.show()

# Custom autograd function example
class CustomReLU(torch.autograd.Function):
    """
    Custom ReLU implementation with autograd support
    """
    @staticmethod
    def forward(ctx, input):
        # Save input for backward pass
        ctx.save_for_backward(input)
        return input.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_output):
        # Retrieve saved input
        input, = ctx.saved_tensors
        # Create gradient mask
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0
        return grad_input

# Test custom ReLU
print("\nTesting custom ReLU function...")
x = torch.randn(3, 3, requires_grad=True)
y = CustomReLU.apply(x)

# Test backward pass
z = y.sum()
z.backward()

print("Input tensor:")
print(x)
print("\nOutput tensor (after custom ReLU):")
print(y)
print("\nGradient of output with respect to input:")
print(x.grad)

# Compare with built-in ReLU
print("\nComparing with built-in ReLU:")
y_builtin = torch.relu(x)
print("Built-in ReLU output:")
print(y_builtin)
print("Difference between custom and built-in ReLU:")
print(torch.abs(y - y_builtin).max().item())

# PyTorch Lightning example - simplified training framework
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping

class LitMNIST(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = Net()
        self.criterion = nn.CrossEntropyLoss()

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        data, target = batch
        output = self(data)
        loss = self.criterion(output, target)
        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        data, target = batch
        output = self(data)
        loss = self.criterion(output, target)
        acc = (output.argmax(dim=1) == target).float().mean()
        self.log('val_loss', loss, prog_bar=True)
        self.log('val_acc', acc, prog_bar=True)
        return loss

    def test_step(self, batch, batch_idx):
        data, target = batch
        output = self(data)
        loss = self.criterion(output, target)
        acc = (output.argmax(dim=1) == target).float().mean()
        self.log('test_loss', loss)
        self.log('test_acc', acc)
        return loss

    def configure_optimizers(self):
        return optim.Adam(self.parameters())

# Initialize model
model = LitMNIST()

# Define callbacks
checkpoint_callback = ModelCheckpoint(
    monitor='val_acc',
    dirpath='./lightning_checkpoints',
    filename='mnist-{epoch:02d}-{val_acc:.2f}',
    save_top_k=3,
    mode='max'
)

early_stopping = EarlyStopping(
    monitor='val_loss',
    patience=3,
    mode='min'
)

# Initialize trainer
trainer = pl.Trainer(
    max_epochs=5,
    callbacks=[checkpoint_callback, early_stopping],
    accelerator='auto',
    devices='auto',
    log_every_n_steps=10
)

# Train the model
print("\nTraining with PyTorch Lightning...")
trainer.fit(model, train_loader, test_loader)

# Test the model
trainer.test(model, dataloaders=test_loader)

print("PyTorch Lightning training completed!")

# Mixed precision training example
from torch.cuda.amp import GradScaler, autocast

# Initialize gradient scaler for mixed precision
scaler = GradScaler()

# Training loop with mixed precision
print("\nMixed precision training example...")
for epoch in range(2):  # Reduced epochs for demonstration
    model.train()
    running_loss = 0.0

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)

        # Forward pass with mixed precision
        with autocast():
            output = model(data)
            loss = criterion(output, target)

        # Backward pass with gradient scaling
        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        running_loss += loss.item()

        if batch_idx % 100 == 99:
            print(f'Epoch: {epoch+1}, Batch: {batch_idx+1}, Loss: {running_loss/100:.3f}')
            running_loss = 0.0

print("Mixed precision training completed!")

# Distributed training example
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP

def setup(rank, world_size):
    """Initialize distributed training"""
    dist.init_process_group(
        backend='nccl',
        init_method='env://',
        world_size=world_size,
        rank=rank
    )

def cleanup():
    """Clean up distributed training"""
    dist.destroy_process_group()

def train(rank, world_size):
    """Training function for distributed training"""
    setup(rank, world_size)

    # Create model and move to GPU
    model = Net().to(rank)
    model = DDP(model, device_ids=[rank])

    # Create optimizer
    optimizer = optim.Adam(model.parameters())

    # Create data loader with distributed sampler
    train_sampler = torch.utils.data.distributed.DistributedSampler(
        train_set,
        num_replicas=world_size,
        rank=rank
    )

    train_loader = torch.utils.data.DataLoader(
        train_set,
        batch_size=64,
        sampler=train_sampler
    )

    # Training loop
    for epoch in range(2):  # Reduced epochs for demonstration
        model.train()
        train_sampler.set_epoch(epoch)
        running_loss = 0.0

        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(rank), target.to(rank)

            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

            if batch_idx % 100 == 99 and rank == 0:
                print(f'Epoch: {epoch+1}, Batch: {batch_idx+1}, Loss: {running_loss/100:.3f}')
                running_loss = 0.0

    cleanup()

# Example usage (would need to be run with torch.distributed.launch)
# world_size = torch.cuda.device_count()
# mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)
print("\nDistributed training setup complete. Use torch.distributed.launch to run.")