Primeros pasos con PyTorch

# CPU only (funciona en cualquier máquina)
pip install torch torchvision

# Con CUDA 12.1 (si tienes GPU NVIDIA)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

print(f"PyTorch version: {torch.__version__}")
print(f"CUDA disponible: {torch.cuda.is_available()}")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Usando device: {device}")

# Transformaciones: convertir a tensor + normalizar
transform = transforms.Compose([
    transforms.ToTensor(),           # PIL → tensor, escala a [0,1]
    transforms.Lambda(lambda x: x.view(-1)),  # (1,28,28) → (784,)
])

# Descargar y crear datasets
train_dataset = datasets.MNIST(root="./data", train=True,
                               download=True, transform=transform)
test_dataset  = datasets.MNIST(root="./data", train=False,
                               download=True, transform=transform)

# Crear DataLoaders
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader  = DataLoader(test_dataset, batch_size=256, shuffle=False)

print(f"Train: {len(train_dataset)} imágenes")
print(f"Test:  {len(test_dataset)} imágenes")
print(f"Batches de train: {len(train_loader)}")

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(784, 128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, 10),
        )

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

# Crear instancia y mover a GPU/CPU
model = MLP().to(device)
print(model)
print(f"\nParámetros totales: {sum(p.numel() for p in model.parameters()):,}")

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

def train_one_epoch(model, loader, criterion, optimizer, device):
    model.train()  # Activa Dropout y BatchNorm en modo entrenamiento
    total_loss = 0
    correct = 0
    total = 0

    for batch_x, batch_y in loader:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)

        # Forward pass
        outputs = model(batch_x)
        loss = criterion(outputs, batch_y)

        # Backward pass
        optimizer.zero_grad()   # Limpiar gradientes anteriores
        loss.backward()         # Calcular gradientes
        optimizer.step()        # Actualizar pesos

        # Estadísticas
        total_loss += loss.item() * batch_x.size(0)
        _, predicted = outputs.max(1)
        correct += predicted.eq(batch_y).sum().item()
        total += batch_y.size(0)

    return total_loss / total, correct / total


# Entrenar 10 epochs
EPOCHS = 10
for epoch in range(EPOCHS):
    train_loss, train_acc = train_one_epoch(
        model, train_loader, criterion, optimizer, device
    )
    print(f"Epoch {epoch+1:2d}/{EPOCHS} | "
          f"Loss: {train_loss:.4f} | Accuracy: {train_acc:.4f}")

def evaluate(model, loader, criterion, device):
    model.eval()  # Desactiva Dropout
    total_loss = 0
    correct = 0
    total = 0

    with torch.no_grad():  # No calcular gradientes → más rápido, menos memoria
        for batch_x, batch_y in loader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            outputs = model(batch_x)
            loss = criterion(outputs, batch_y)

            total_loss += loss.item() * batch_x.size(0)
            _, predicted = outputs.max(1)
            correct += predicted.eq(batch_y).sum().item()
            total += batch_y.size(0)

    return total_loss / total, correct / total


test_loss, test_acc = evaluate(model, test_loader, criterion, device)
print(f"🎯 Test Loss: {test_loss:.4f}")
print(f"🎯 Test Accuracy: {test_acc:.4f}")
print(f"Acierta {test_acc*100:.1f}% de los dígitos que nunca ha visto")

# Opción 1 (RECOMENDADA): guardar solo los pesos
torch.save(model.state_dict(), "mi_modelo_mnist.pth")
print("✅ Pesos guardados como mi_modelo_mnist.pth")

# Opción 2: guardar todo (modelo + pesos + optimizer)
checkpoint = {
    "model_state": model.state_dict(),
    "optimizer_state": optimizer.state_dict(),
    "epoch": EPOCHS,
    "test_acc": test_acc,
}
torch.save(checkpoint, "checkpoint_mnist.pth")
print("✅ Checkpoint completo guardado")

# En un nuevo script o sesión:
import torch
import torch.nn as nn

# 1. Recrear la MISMA arquitectura
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(784, 128), nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(128, 64),  nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(64, 10),
        )
    def forward(self, x):
        return self.model(x)

# 2. Crear instancia y cargar pesos
loaded_model = MLP()
loaded_model.load_state_dict(torch.load("mi_modelo_mnist.pth", weights_only=True))
loaded_model.eval()  # Modo evaluación
print("✅ Modelo cargado correctamente")

# Tomar 5 imágenes del test set
test_images, test_labels = [], []
for img, label in test_dataset:
    test_images.append(img)
    test_labels.append(label)
    if len(test_images) == 5:
        break

# Stack en un batch (5, 784)
batch = torch.stack(test_images)

# Predecir
with torch.no_grad():
    logits = loaded_model(batch)
    probs  = torch.softmax(logits, dim=1)
    preds  = probs.argmax(dim=1)

# Mostrar resultados
for i in range(5):
    confidence = probs[i][preds[i]].item() * 100
    real = test_labels[i]
    status = "✅" if preds[i].item() == real else "❌"
    print(f"{status} Imagen {i}: predicho={preds[i].item()} "
          f"(confianza: {confidence:.1f}%) | real={real}")