File size: 4,237 Bytes
1fff313 5144b79 55af1cd f774571 55af1cd 1fff313 f774571 55af1cd 1fff313 5144b79 1fff313 f774571 1fff313 f774571 1fff313 5144b79 1fff313 f774571 1fff313 f774571 1fff313 f774571 1fff313 f774571 1fff313 91c78a9 5144b79 91c78a9 1fff313 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.tensorboard.writer import SummaryWriter
from mnist_classifier.dataset import MNISTDataModule
from mnist_classifier.model import MNISTModel
from datetime import datetime
import os
import random
import numpy as np
from tqdm import tqdm
def set_seed(seed):
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
random.seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
def train():
# Training loop
learning_rate = 0.001
batch_size = 128
epochs = 10
# Set seed for reproducibility
set_seed(42)
# Set device
device = torch.device('cuda')
print(f"Using device: {device}")
# Initialize tensorboard
log_dir = 'runs/mnist_experiment_' + datetime.now().strftime('%Y%m%d-%H%M%S')
writer = SummaryWriter(log_dir)
# Setup data
data_module = MNISTDataModule(batch_size=batch_size, val_batch_size=1000)
train_loader, test_loader = data_module.get_dataloaders()
# Initialize model, optimizer, and loss function
model = MNISTModel().to(device)
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
criterion = nn.CrossEntropyLoss()
num_epochs = epochs
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
correct = 0
total = 0
with tqdm(total=len(train_loader), desc=f"Epoch {epoch+1}/{num_epochs}", unit="batch") as pbar:
for batch_idx, batch in enumerate(train_loader):
images, labels = batch[0].to(device), batch[1].to(device)
if batch_idx == 0:
print(f"images shape: {images.shape}")
print(f"labels shape: {labels.shape}")
# print number of images in batch
print(f"Number of images in batch: {len(images)}")
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
# Update tqdm progress bar
pbar.set_postfix({
'loss': running_loss / (batch_idx + 1),
'accuracy': 100. * correct / total,
'step': batch_idx + 1
})
pbar.update(1)
if batch_idx % 100 == 99:
writer.add_scalar('training loss',
running_loss / 100,
epoch * len(train_loader) + batch_idx)
writer.add_scalar('training accuracy',
100. * correct / total,
epoch * len(train_loader) + batch_idx)
running_loss = 0.0
# Validation phase
model.eval()
test_loss = 0
correct = 0
total = 0
with torch.no_grad():
for batch in test_loader:
images = batch[0].to(device)
labels = batch[1].to(device)
outputs = model(images)
loss = criterion(outputs, labels)
test_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
accuracy = 100. * correct / total
writer.add_scalar('test accuracy', accuracy, epoch)
print(f'Epoch {epoch+1}: Test Accuracy: {accuracy:.2f}%')
writer.close()
# Ensure the directory exists
os.makedirs("./models", exist_ok=True)
# Format the filename with the config parameters
filename = f"./models/mnist_model_lr{learning_rate}_bs{batch_size}_ep{epochs}.pth"
torch.save(model.state_dict(), filename)
if __name__ == "__main__":
train() |