Spaces:
Sleeping
Sleeping
File size: 5,300 Bytes
068b6e0 | 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 128 129 130 131 132 | import os
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, models
import medmnist
from medmnist import INFO
from torch.utils.data import DataLoader
from tqdm import tqdm
import matplotlib.pyplot as plt
def main():
# 1. Setup and Hardware Configuration
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Training on: {device}")
# Set this to point to your secondary NVMe drive to prevent OS drive I/O bottlenecks
dataset_root = r"C:\Users\USER\Downloads\MedMNIST_Data"
os.makedirs(dataset_root, exist_ok=True)
data_flag = 'pneumoniamnist'
info = INFO[data_flag]
DataClass = getattr(medmnist, info['python_class'])
# 2. The Golden Preprocessing & Dynamic Augmentation
# We normalize to [-1, 1] using mean=0.5, std=0.5 so it matches your team's generator math
train_transform = transforms.Compose([
transforms.Grayscale(num_output_channels=3), # ResNet expects 3 RGB channels
transforms.RandomHorizontalFlip(), # Dynamic spatial augmentation
transforms.RandomRotation(10), # Dynamic spatial augmentation
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
val_transform = transforms.Compose([
transforms.Grayscale(num_output_channels=3),
transforms.ToTensor(), # NO spatial augmentation for validation
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
# 3. Load Datasets
print("Fetching 224x224 dataset...")
train_dataset = DataClass(split='train', transform=train_transform, download=True, size=224, root=dataset_root)
val_dataset = DataClass(split='val', transform=val_transform, download=True, size=224, root=dataset_root)
# 4. DataLoaders
# Using batch size 32. num_workers=0 is the safest default for Windows to prevent multiprocessing crashes.
train_loader = DataLoader(dataset=train_dataset, batch_size=32, shuffle=True, num_workers=0)
val_loader = DataLoader(dataset=val_dataset, batch_size=32, shuffle=False, num_workers=0)
# 5. Initialize ResNet50
print("Loading ResNet50...")
model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
# Modify the final layer for Binary Classification (Pneumonia vs Normal)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 2)
model = model.to(device)
# 6. Loss and Optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-4) # 1e-4 is a very stable learning rate for fine-tuning
num_epochs = 10
history_loss = []
history_acc = []
# 7. The Training Loop
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
correct = 0
total = 0
# tqdm creates a nice progress bar in the terminal
loop = tqdm(train_loader, leave=True)
for images, labels in loop:
images, labels = images.to(device), labels.to(device).squeeze().long()
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
loop.set_description(f"Epoch [{epoch+1}/{num_epochs}]")
loop.set_postfix(loss=loss.item(), acc=100.*correct/total)
# Calculate the average loss and accuracy for this epoch
epoch_loss = running_loss / len(train_loader)
epoch_acc = 100. * correct / total
history_loss.append(epoch_loss)
history_acc.append(epoch_acc)
# 8. Save the Frozen Weights for your team
save_path = os.path.join(dataset_root, 'baseline_resnet50.pth')
torch.save(model.state_dict(), save_path)
print(f"\nTraining Complete! Baseline weights saved to: {save_path}")
# Create the Learning Curve Graph
fig, ax1 = plt.subplots(figsize=(10, 6))
# Plot Loss (Red Line)
color = 'tab:red'
ax1.set_xlabel('Epochs', fontweight='bold')
ax1.set_ylabel('Training Loss', color=color, fontweight='bold')
ax1.plot(range(1, num_epochs+1), history_loss, color=color, marker='o', label='Loss')
ax1.tick_params(axis='y', labelcolor=color)
# Plot Accuracy (Blue Line) on the same graph
ax2 = ax1.twinx()
color = 'tab:blue'
ax2.set_ylabel('Training Accuracy (%)', color=color, fontweight='bold')
ax2.plot(range(1, num_epochs+1), history_acc, color=color, marker='s', label='Accuracy')
ax2.tick_params(axis='y', labelcolor=color)
plt.title('ResNet50 Training Curve', fontsize=14, fontweight='bold')
fig.tight_layout()
# Save the image
graph_path = os.path.join(dataset_root, 'learning_curve.png')
plt.savefig(graph_path, dpi=300)
print(f"Learning Curve saved to: {graph_path}")
if __name__ == '__main__':
main() |