============================================================

COMPARAÇÃO DE BATCH SIZE NO TREINAMENTO

============================================================

import time
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

------------------------------------------------------------

Exemplo de hiperparâmetros-base

------------------------------------------------------------

base_args = {
'num_workers': 4,
'num_classes': 10,
'lr': 1e-4,
'weight_decay': 5e-4,
'num_epochs': 30,
'device': 'cuda' if torch.cuda.is_available() else 'cpu'
}

------------------------------------------------------------

Função de treino por 1 época

------------------------------------------------------------

def train_one_epoch(loader, net, criterion, optimizer, device):
net.train()
losses = []

for dado, rotulo in loader:
    dado = dado.to(device)
    rotulo = rotulo.to(device)

    pred = net(dado)
    loss = criterion(pred, rotulo)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    losses.append(loss.item())

return np.mean(losses), np.std(losses)

------------------------------------------------------------

Função principal para testar um batch_size

------------------------------------------------------------

def run_experiment(batch_size, train_dataset, build_model):
args = base_args.copy()
args['batch_size'] = batch_size

# DataLoader
train_loader = DataLoader(
    train_dataset,
    batch_size=args['batch_size'],
    shuffle=True,
    num_workers=args['num_workers']
)

# Modelo novo para cada experimento
net = build_model().to(args['device'])

# Loss e otimizador
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(
    net.parameters(),
    lr=args['lr'],
    weight_decay=args['weight_decay']
)

epoch_times = []
epoch_below_16 = None
history = []

print(f"\n===== Testando batch_size = {batch_size} =====")

for epoch in range(1, args['num_epochs'] + 1):
    start = time.time()

    mean_loss, std_loss = train_one_epoch(
        train_loader, net, criterion, optimizer, args['device']
    )

    elapsed = time.time() - start
    epoch_times.append(elapsed)
    history.append(mean_loss)

    print(
        f"Época {epoch:02d} | "
        f"Loss: {mean_loss:.4f} +- {std_loss:.4f} | "
        f"Tempo: {elapsed:.2f}s"
    )

    # Registra a primeira época em que loss < 1.6
    if epoch_below_16 is None and mean_loss < 1.6:
        epoch_below_16 = epoch

avg_epoch_time = np.mean(epoch_times)

return {
    'batch_size': batch_size,
    'avg_epoch_time': avg_epoch_time,
    'epoch_below_16': epoch_below_16,
    'loss_history': history
}

------------------------------------------------------------

Exemplo de função que constrói a rede

Ajuste conforme o notebook da aula

------------------------------------------------------------

def build_model():
return nn.Sequential(
nn.Flatten(),
nn.Linear(28 * 28, 128),
nn.ReLU(),
nn.Linear(128, 10)
)

------------------------------------------------------------

EXECUÇÃO DOS EXPERIMENTOS

train_dataset deve já existir no notebook

------------------------------------------------------------

results = []

for bs in [100, 20, 5]:
result = run_experiment(bs, train_dataset, build_model)
results.append(result)

------------------------------------------------------------

RELATÓRIO FINAL

------------------------------------------------------------

print("\n========== RESUMO FINAL ==========")
for r in results:
print(f"\nBatch size: {r['batch_size']}")
print(f"Tempo médio por época: {r['avg_epoch_time']:.2f} segundos")

if r['epoch_below_16'] is not None:
    print(f"Loss < 1.6 na época: {r['epoch_below_16']}")
else:
    print("Não atingiu loss < 1.6 em 30 épocas")