import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, random_split
from tqdm import tqdm

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

transforms_classic = transforms.Compose([
    transforms.Pad(2),
    transforms.ToTensor(),
    transforms.Normalize(0.5, 0.5)
])

train_dataset = datasets.FashionMNIST(root='./data', train=True, transform=transforms_classic, download=True)
test_dataset = datasets.FashionMNIST(root='./data', train=False, transform=transforms_classic, download=True)

train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4, pin_memory=True, persistent_workers=True)
val_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=4, pin_memory=True, persistent_workers=True)

100%|██████████| 26.4M/26.4M [00:01<00:00, 13.7MB/s]
100%|██████████| 29.5k/29.5k [00:00<00:00, 201kB/s]
100%|██████████| 4.42M/4.42M [00:01<00:00, 3.74MB/s]
100%|██████████| 5.15k/5.15k [00:00<00:00, 18.8MB/s]

class LeNetClassic(nn.Module):
  def __init__(self):
    super(LeNetClassic, self).__init__()
    self.features = nn.Sequential(
        nn.Conv2d(1, 6, kernel_size=5, stride=1),
        nn.Tanh(),
        nn.AvgPool2d(kernel_size=2, stride=2),
        nn.Conv2d(6, 16, kernel_size=5, stride=1),
        nn.Tanh(),
        nn.AvgPool2d(kernel_size=2, stride=2)
    )

    self.classfier = nn.Sequential(
        nn.Linear(16 * 5 * 5, 120),
        nn.Tanh(),
        nn.Linear(120, 84),
        nn.Tanh(),
        nn.Linear(84, 10)
    )

  def forward(self, x):
    x = self.features(x)
    x = torch.flatten(x, 1)
    x = self.classfier(x)
    return x

model = LeNetClassic().to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)

best_acc = 0
epochs = 50

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

  train_pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs} [Train]")
  for inputs, targets in train_pbar:
    inputs, targets = inputs.to(device), targets.to(device)
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()

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

    train_pbar.set_postfix(loss=running_loss/(total/targets.size(0)), acc=100.*correct/total)

  train_loss = running_loss / len(train_loader)
  train_acc = 100. * correct / total

  model.eval()
  val_loss = 0.0
  val_correct = 0
  val_total = 0

  with torch.no_grad():
    val_pbar = tqdm(val_loader, desc=f"Epoch {epoch+1}/{epochs} [Val]")
    for inputs, targets in val_pbar:
      inputs, targets = inputs.to(device), targets.to(device)
      outputs = model(inputs)
      loss = criterion(outputs, targets)

      val_loss += loss.item()
      _, predicted = outputs.max(1)
      val_total += targets.size(0)
      val_correct += predicted.eq(targets).sum().item()
      val_pbar.set_postfix(loss=val_loss/(val_total/targets.size(0)), acc=100.*val_correct/val_total)

  validation_loss = val_loss / len(val_loader)
  validation_acc = 100. * val_correct / val_total

  print(f"Epoch {epoch+1}/{epochs} | Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}% | Val Loss: {validation_loss:.4f}, Val Acc: {validation_acc:.2f}%")

  scheduler.step()

  if validation_acc > best_acc:
    best_acc = validation_acc
    torch.save(model.state_dict(), 'best_lenet_classic_model.pth')
    print(f"New best model found! Accuracy: {best_acc:.2f}%. Saving model...")

print("\nTraining finished.")
print(f"Best validation accuracy: {best_acc:.2f}%")

Epoch 1/50 [Train]: 100%|██████████| 469/469 [00:05<00:00, 85.56it/s, acc=68, loss=0.671] 
Epoch 1/50 [Val]: 100%|██████████| 79/79 [00:00<00:00, 93.31it/s, acc=78.1, loss=0.074]

Epoch 1/50 | Train Loss: 0.8939, Train Acc: 68.05% | Val Loss: 0.5856, Val Acc: 78.07%
New best model found! Accuracy: 78.07%. Saving model...

Epoch 2/50 [Train]: 100%|██████████| 469/469 [00:04<00:00, 116.00it/s, acc=81.5, loss=0.378]
Epoch 2/50 [Val]: 100%|██████████| 79/79 [00:00<00:00, 112.89it/s, acc=82.5, loss=0.0606]

Epoch 2/50 | Train Loss: 0.5036, Train Acc: 81.51% | Val Loss: 0.4797, Val Acc: 82.55%
New best model found! Accuracy: 82.55%. Saving model...

Epoch 3/50 [Train]: 100%|██████████| 469/469 [00:03<00:00, 118.70it/s, acc=84.6, loss=0.319]
Epoch 3/50 [Val]: 100%|██████████| 79/79 [00:00<00:00, 116.32it/s, acc=84.1, loss=0.0547]

Epoch 3/50 | Train Loss: 0.4252, Train Acc: 84.55% | Val Loss: 0.4328, Val Acc: 84.13%
New best model found! Accuracy: 84.13%. Saving model...

Epoch 4/50 [Train]: 100%|██████████| 469/469 [00:04<00:00, 111.34it/s, acc=86, loss=0.287]
Epoch 4/50 [Val]: 100%|██████████| 79/79 [00:00<00:00, 115.16it/s, acc=86.3, loss=0.0489]

Epoch 4/50 | Train Loss: 0.3824, Train Acc: 86.00% | Val Loss: 0.3872, Val Acc: 86.29%
New best model found! Accuracy: 86.29%. Saving model...

Epoch 5/50 [Train]: 100%|██████████| 469/469 [00:04<00:00, 117.14it/s, acc=87.2, loss=0.264]
Epoch 5/50 [Val]: 100%|██████████| 79/79 [00:00<00:00, 115.73it/s, acc=87, loss=0.0464]

Epoch 5/50 | Train Loss: 0.3523, Train Acc: 87.22% | Val Loss: 0.3672, Val Acc: 86.96%
New best model found! Accuracy: 86.96%. Saving model...