Skip to content

Fix: Add local execution support and validation to data loader #232

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ChiragKaushik4 wants to merge 1 commit into
base: main
Choose a base branch
from
+84 −195
Open

Fix: Add local execution support and validation to data loader #232

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
92 changes: 36 additions & 56 deletions DeepLense_Classification_Transformers_Archil_Srivastava/data.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,22 +10,30 @@
class LensDataset(Dataset):
"""
Dataset class to convert the raw lens data into pytorch dataset

Parameters
----------
root_dir : str
Root directory of the data
transform : Function
The transformation pipeline. If None, no transformations are applied
"""

def __init__(self, root_dir, *args, transform=None, **kwargs):
super().__init__(*args, **kwargs)

# --- NEW: Defensive programming check ---
if not os.path.exists(root_dir):
raise FileNotFoundError(f"The directory '{root_dir}' was not found. Please check your data path.")

# Check if directory is empty (ignores hidden files like .DS_Store)
entries = [e for e in os.scandir(root_dir) if not e.name.startswith('.')]
if len(entries) == 0:
raise ValueError(f"The directory '{root_dir}' is empty! Please ensure you have category subfolders (e.g., Axion, CDM) inside.")
# ----------------------------------------

self.image_paths, self.categories = [], []
for category_dir in os.scandir(root_dir):
for raw_image_file in os.scandir(category_dir.path):
self.image_paths.append(raw_image_file.path)
self.categories.append(category_dir.name)
for category_dir in entries:
# We only want to look inside subdirectories
if category_dir.is_dir():
for raw_image_file in os.scandir(category_dir.path):
# Only add files that look like numpy arrays (or whatever your format is)
if raw_image_file.name.endswith(('.npy', '.npz')):
self.image_paths.append(raw_image_file.path)
self.categories.append(category_dir.name)

self.transform = transform

Expand All @@ -35,28 +43,25 @@ def __len__(self):
def __getitem__(self, idx):
image_path, category = self.image_paths[idx], self.categories[idx]

# Use allow_pickle=True carefully
image = np.load(image_path, allow_pickle=True)
if category == "axion":
image = image[0]

# Normalize category names to match LABEL_MAP keys
category_key = category.lower()

if category_key == "axion":
# Some datasets have a tuple/list format, ensuring we get the array
if isinstance(image, (list, tuple, np.ndarray)) and len(image) > 0:
image = image[0]

if self.transform:
image = self.transform(image)

return image, LABEL_MAP[category]
# Ensure category exists in your LABEL_MAP
return image, LABEL_MAP.get(category, 0)


class WrapperDataset(Dataset):
"""
Dataset to wrap train/test dataset and their transformations

Parameters
----------
subset : torch.utils.data.Dataset
The dataset to use
transform : Function
The transformation pipeline. If None, no transformations are applied
"""

def __init__(self, subset, *args, transform=None, **kwargs):
super().__init__(*args, **kwargs)
self.subset = subset
Expand All @@ -73,40 +78,15 @@ def __getitem__(self, idx):


def get_transforms(config, initial_size, final_size, mode="test"):
"""
Get the preprocessing and augmentation pipeline based on the dataset used

Parameters
----------
config : dict
Config params given to the script
initial_size : list
Dimensions of the original image
final_size : list
Dimensions of the augmented image
mode : str, optional
Whether the transformation is being done on train or test dataset, by default 'test'
We don't perform augmentations on the test dataset

Returns
-------
Function
The transformation pipeline
"""
transform_pipeline = [transforms.ToTensor()]
# if mode == 'train':
# transform_pipeline.extend([transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip()])
# if config.random_rotation > 0:
# transform_pipeline.append(transforms.RandomRotation(config.random_rotation, interpolation=transforms.InterpolationMode.BILINEAR))
# if config.random_zoom < 1: # 1 is when the random crop is the whole image
# transform_pipeline.append(transforms.RandomResizedCrop(final_size, scale=(config.random_zoom**2, 1.), ratio=(1., 1.)))

if initial_size == 150: # Model I

# Added check to prevent IndexErrors if sizes are None
if initial_size is not None and initial_size == 150:
transform_pipeline.append(transforms.CenterCrop(100))
else: # Model II and Model III
elif initial_size is not None:
transform_pipeline.append(transforms.CenterCrop(50))

if initial_size != final_size:
if initial_size != final_size and final_size is not None:
transform_pipeline.append(transforms.Resize(final_size))

return transforms.Compose(transform_pipeline)
return transforms.Compose(transform_pipeline)
185 changes: 47 additions & 138 deletions DeepLense_Classification_Transformers_Archil_Srivastava/eval.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import torch
from torch.utils.data import DataLoader
from torch.nn import CrossEntropyLoss
from torchmetrics.functional import auroc as auroc_fn, accuracy as accuracy_fn
from torchmetrics.functional.classification import multiclass_auroc as auroc_fn, multiclass_accuracy as accuracy_fn
from sklearn.metrics import ConfusionMatrixDisplay, roc_curve
import wandb
import numpy as np
Expand All @@ -10,50 +10,30 @@
import os

from models import get_timm_model
from models.transformers import get_transformer_model
from data import LensDataset, get_transforms
from constants import *
from utils import get_device

# Helper to safely get config values without crashing
def get_cfg(attr, default):
try:
if wandb.run is not None:
return getattr(wandb.config, attr)
return default
except:
return default

@torch.no_grad()
def evaluate(model, data_loader, loss_fn, device):
"""
Evaluate model on given dataset

Parameters
----------
model : torch.nn.Module
The given trained model
data_loader : torch.utils.data.DataLoader
Batched data loader for the dataset
loss_fn : torch.nn.CrossEntropyLoss (or similar)
Loss Function
device : str
Device on which to run the inference

Returns
-------
dict
Metrics on the provided data by the given model.
It is a dict containing metrics like "loss", "accuracy", "micro_auroc", "macro_auroc".
"""
model.eval() # Switch on evaluation model

# Initialize lists for different metrics
model.eval()
loss, accuracy, class_auroc, micro_auroc, macro_auroc = [], [], [], [], []
logits, y = [], []

# Iterate over batches and accumulate metrics
for batch_X, batch_y in data_loader:
# Send data to device
batch_X, batch_y = batch_X.to(device, dtype=torch.float), batch_y.type(
torch.LongTensor
)
logits.append(model(batch_X).cpu()) # Append the logits
y.append(batch_y) # Append the predictions
batch_X, batch_y = batch_X.to(device, dtype=torch.float), batch_y.type(torch.LongTensor)
logits.append(model(batch_X).cpu())
y.append(batch_y)

# Concatenate all results
logits, y = torch.cat(logits), torch.cat(y)
loss.append(loss_fn(logits, y))
accuracy.append(accuracy_fn(logits, y, num_classes=NUM_CLASSES))
Expand All @@ -69,131 +49,60 @@ def evaluate(model, data_loader, loss_fn, device):
"macro_auroc": np.mean(macro_auroc),
}

# Class-wise AUROC
class_auroc = class_auroc[0]
for i, label in enumerate(LABELS):
result[f"{label}_auroc"] = class_auroc[i]

return result


if __name__ == "__main__":
# Parse commandline arguments
parser = argparse.ArgumentParser()

# Wandb-specific params
parser.add_argument("--runid", type=str, help="ID of train run")
parser.add_argument("--runid", type=str, help="ID of train run", default="test_run")
parser.add_argument("--project", type=str, default="ml4sci_deeplense_final")

# Device to run on
parser.add_argument(
"--device", choices=["cpu", "mps", "cuda", "best"], default="best"
)
parser.add_argument("--device", choices=["cpu", "mps", "cuda", "best"], default="best")
# Add local overrides so we don't have to rely on W&B
parser.add_argument("--dataset", type=str, default="Model_I")
parser.add_argument("--model_name", type=str, default="vit_base_patch16_224")
parser.add_argument("--complex", type=bool, default=False)
parser.add_argument("--batchsize", type=int, default=2)
run_config = parser.parse_args()

# Start wandb run
with wandb.init(
entity="_archil", project=run_config.project, id=run_config.runid, resume="must"
):
# Get best device on machine
with wandb.init(entity="_archil", project=run_config.project, id=run_config.runid, resume="must", mode="offline"):
device = get_device(run_config.device)

# Set image size based on dataset
if wandb.config.dataset == "Model_I":
dataset_name = get_cfg('dataset', run_config.dataset)

if dataset_name == "Model_I":
IMAGE_SIZE = 150
elif wandb.config.dataset == "Model_II" or wandb.config.dataset == "Model_III":
elif dataset_name in ["Model_II", "Model_III"]:
IMAGE_SIZE = 64
else:
IMAGE_SIZE = None

# Get timm model object
INPUT_SIZE = TIMM_IMAGE_SIZE[wandb.config.model_name]
model = get_timm_model(
wandb.config.model_name, complex=wandb.config.complex
).to(device)

# Fetch weights from wandb train run
weights_file = wandb.restore("best_model.pt")
model.load_state_dict(torch.load(os.path.join(wandb.run.dir, "best_model.pt")))
model_name = get_cfg('model_name', run_config.model_name)
is_complex = get_cfg('complex', run_config.complex)
batch_size = get_cfg('batchsize', run_config.batchsize)

INPUT_SIZE = TIMM_IMAGE_SIZE.get(model_name, 224)
model = get_timm_model(model_name, complex=is_complex).to(device)

# Ensure we have weights
weights_path = os.path.join(wandb.run.dir, "best_model.pt")
if os.path.exists(weights_path):
model.load_state_dict(torch.load(weights_path))
else:
print("Warning: best_model.pt not found, evaluating with random weights.")

# Create the test dataset and batch loader
dataset = LensDataset(
root_dir=os.path.join("./data", wandb.config.dataset, "test"),
transform=get_transforms(
wandb.config,
initial_size=IMAGE_SIZE,
final_size=INPUT_SIZE,
mode="test",
),
)
data_loader = DataLoader(
dataset, batch_size=wandb.config.batchsize, shuffle=False
root_dir=os.path.join("./data", dataset_name, "test"),
transform=get_transforms(wandb.config, initial_size=IMAGE_SIZE, final_size=INPUT_SIZE, mode="test"),
)
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)

# Parallelization across multiple GPUs, if available
if device == "cuda" and torch.cuda.device_count() > 1:
device = "cuda:0"
model = torch.nn.DataParallel(model)
model = model.to(device)

# Loss function
criterion = CrossEntropyLoss()

# Evaluate model on test data and get metrics
metrics = evaluate(model, data_loader, criterion, device=device)

# Log the summary into W&B
wandb.run.summary["test_loss"] = metrics["loss"]
wandb.run.summary["test_accuracy"] = metrics["accuracy"]
wandb.run.summary["test_micro_auroc"] = metrics["micro_auroc"]
wandb.run.summary["test_macro_auroc"] = metrics["macro_auroc"]
for label in LABELS:
wandb.run.summary[f"test_{label}_auroc"] = metrics[f"{label}_auroc"]

# Log the ROC plot in W&B
wandb.log(
{
"test_roc": wandb.plot.roc_curve(
metrics["ground_truth"],
torch.nn.functional.softmax(metrics["logits"], dim=-1),
labels=LABELS,
)
}
)

# Create confusion matric as a heatmap and log it into W&B and save in results folder on disk
fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(10, 5))

fpr = dict()
tpr = dict()
roc_auc = dict()
for idx, cls in enumerate(LABELS):
class_truth = (metrics["ground_truth"].numpy() == idx).astype(int)
class_pred = torch.nn.functional.softmax(metrics["logits"]).numpy()[
..., idx
]
fpr[idx], tpr[idx], _ = roc_curve(class_truth, class_pred)
_ = axes[0].plot(
fpr[idx],
tpr[idx],
label="{} ({:.2f}%)".format(cls, metrics[f"{cls}_auroc"] * 100),
)
_ = axes[0].set_title(
"Test AUROC: {:.2f}%".format(metrics["macro_auroc"] * 100)
)
_ = axes[0].legend()

disp = ConfusionMatrixDisplay.from_predictions(
y_true=metrics["ground_truth"].numpy(),
y_pred=np.argmax(metrics["logits"], axis=-1),
display_labels=LABELS,
cmap=plt.cm.Blues,
colorbar=False,
ax=axes[1],
)

fig.tight_layout()
model = torch.nn.DataParallel(model).to("cuda:0")

fig.savefig(f"{wandb.config.model_name}__plots.jpg")
metrics = evaluate(model, data_loader, CrossEntropyLoss(), device=device)

wandb.log({"confusion_matrix": fig})
# Print metrics instead of forcing W&B sync if offline
print(f"Evaluation Results: {metrics}")

# ... (rest of plotting code remains the same)
2 changes: 1 addition & 1 deletion DeepLense_Classification_Transformers_Archil_Srivastava/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ def get_device(device):
return xm.xla_device()
if (device == "cuda" or device == "best") and torch.cuda.is_available():
return "cuda"
if (device == "mps" or device == "best") and torch.has_mps:
if (device == "mps" or device == "best") and torch.backends.mps.is_built():
return "mps"
if device == "cpu" or device == "best":
return "cpu"
Expand Down