Spaces:
Sleeping
Sleeping
File size: 3,583 Bytes
9916246 | 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 | import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm
from dataset import VinDrCXRClassificationDataset, get_train_val_transforms
from model import CliniScanClassifier
from sklearn.metrics import roc_auc_score
def train_one_epoch(model, loader, optimizer, criterion, device):
model.train()
running_loss = 0.0
loop = tqdm(loader, desc="Training Epoch")
for images, targets in loop:
images, targets = images.to(device), targets.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
running_loss += loss.item() * images.size(0)
loop.set_postfix(loss=loss.item())
return running_loss / len(loader.dataset)
@torch.no_grad()
def evaluate(model, loader, criterion, device):
model.eval()
running_loss = 0.0
all_targets = []
all_outputs = []
for images, targets in tqdm(loader, desc="Evaluating"):
images, targets = images.to(device), targets.to(device)
outputs = model(images)
loss = criterion(outputs, targets)
running_loss += loss.item() * images.size(0)
all_targets.append(targets.cpu())
all_outputs.append(torch.sigmoid(outputs).cpu())
all_targets = torch.cat(all_targets).numpy()
all_outputs = torch.cat(all_outputs).numpy()
# Calculate Macro AUC
try:
auc = roc_auc_score(all_targets, all_outputs, average='macro', multi_class='ovr')
except ValueError:
auc = 0.0 # handles edge case where batch has only one class
return running_loss / len(loader.dataset), auc
def main():
if torch.backends.mps.is_available():
device = torch.device('mps')
elif torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
print(f"Using device: {device}")
# Dummy paths for local
train_csv = './data/train.csv'
val_csv = './data/train.csv' # For demonstration, reusing same dataset for val
img_dir = './data/images'
train_transform, val_transform = get_train_val_transforms()
try:
train_dataset = VinDrCXRClassificationDataset(train_csv, img_dir, transform=train_transform)
val_dataset = VinDrCXRClassificationDataset(val_csv, img_dir, transform=val_transform)
except FileNotFoundError:
print("Data files not found. Skipping dataset initialization for demo purposes.")
return
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)
model = CliniScanClassifier(num_classes=15).to(device)
# Using BCEWithLogitsLoss for Multi-Label Classification
criterion = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
epochs = 10
best_auc = 0.0
for epoch in range(epochs):
train_loss = train_one_epoch(model, train_loader, optimizer, criterion, device)
val_loss, val_auc = evaluate(model, val_loader, criterion, device)
print(f"Epoch {epoch+1}/{epochs} - Train Loss: {train_loss:.4f} - Val Loss: {val_loss:.4f} - Val AUC: {val_auc:.4f}")
if val_auc > best_auc:
best_auc = val_auc
torch.save(model.state_dict(), "best_resnet_classification.pth")
print("Saved Best Model!")
if __name__ == '__main__':
main()
|