Spaces:
Sleeping
Sleeping
File size: 3,101 Bytes
0b86da8 | 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 | import os
import sys
from pathlib import Path
# Add project root to sys.path
sys.path.append(str(Path(__file__).parent.parent))
import matplotlib # noqa: E402
matplotlib.use("Agg") # Use headless backend
import matplotlib.pyplot as plt # noqa: E402
import torch # noqa: E402
import torch.nn as nn # noqa: E402
import yaml # noqa: E402
from sklearn.metrics import ( # noqa: E402
ConfusionMatrixDisplay,
classification_report,
confusion_matrix,
)
def load_config(config_path="config.yaml"):
with open(config_path, "r") as f:
return yaml.safe_load(f)
config = load_config()
CLASSES = config["classes"]
def get_device(config_device):
if config_device == "auto":
return "cuda" if torch.cuda.is_available() else "cpu"
return config_device
DEVICE = get_device(config["device"])
def evaluate(model, data_loader, device=DEVICE, save_dir="models/plots"):
"""
Evaluates a PyTorch model on a given DataLoader.
Args:
model: The PyTorch model to evaluate.
data_loader: The DataLoader providing the evaluation data.
device: The device to run evaluation on (e.g., 'cuda', 'cpu').
save_dir: Directory to save plots.
Returns:
avg_loss (float): The average loss over the dataset.
accuracy (float): The classification accuracy (0.0 to 1.0).
"""
model.to(device)
model.eval()
criterion = nn.CrossEntropyLoss()
total_loss = 0.0
correct = 0
total = 0
all_preds = []
all_labels = []
with torch.no_grad():
for images, labels in data_loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
total_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
all_preds.extend(predicted.cpu().numpy())
all_labels.extend(labels.cpu().numpy())
avg_loss = total_loss / len(data_loader)
accuracy = correct / total
print("\nEvaluation Results:")
print(f"Average Loss: {avg_loss:.4f}")
print(f"Accuracy: {accuracy:.4f}")
# Classification Report
print("\nClassification Report:")
report = classification_report(
all_labels, all_preds, target_names=CLASSES, labels=range(len(CLASSES)), zero_division=0
)
print(report)
# Confusion Matrix
cm = confusion_matrix(all_labels, all_preds, labels=range(len(CLASSES)))
os.makedirs(save_dir, exist_ok=True)
fig, ax = plt.subplots(figsize=(10, 8))
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=CLASSES)
disp.plot(cmap=plt.cm.Blues, ax=ax, xticks_rotation=45)
plt.title("Confusion Matrix")
plt.tight_layout()
plt.savefig(f"{save_dir}/confusion_matrix.png")
print(f"\nConfusion matrix saved to {save_dir}/confusion_matrix.png")
return avg_loss, accuracy
|