tml26_task2 / task_template.py
maitri01's picture
Update task_template.py
8bbcd1d verified
import os
import sys
import requests
from pathlib import Path
import torch
import torch.nn as nn
from torchvision import datasets, transforms
from torchvision.models import resnet18
from safetensors.torch import load_file
import pandas as pd
# --------------------------------
# LOADING A MODEL (EXAMPLE: TARGET MODEL)
# --------------------------------
def make_model():
model = resnet18(weights=None)
model.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
model.maxpool = nn.Identity()
model.fc = nn.Linear(model.fc.in_features, 100)
return model
checkpoint_path = "path/to/your/model_checkpoint.safetensors" # Replace with your model checkpoint path
state_dict = load_file(checkpoint_path, device="cpu")
model = make_model()
model.load_state_dict(state_dict, strict=True)
model.eval()
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5071, 0.4867, 0.4408),
(0.2675, 0.2565, 0.2761)),
])
data_root = "path/to/cifar100" # Replace with your CIFAR-100 dataset path, or where it should be downloaded
dataset = datasets.CIFAR100(root=data_root, train=False, download=True, transform=transform)
x, y = dataset[0] # Example: get the first image and label
with torch.no_grad():
logits = model(x.unsqueeze(0))
print("True label:", y)
print("Logits shape:", logits.shape) # Should be [1, 100] for CIFAR-100
print("Logits:", logits)
# # --------------------------------
# # SUBMISSION FORMAT
# # --------------------------------
"""
The submission must be a .csv file with the following format:
-"id": ID of the subset (from 0 to 359)
-"score": Stealing confidence score for each model (float)
"""
# Example Submission:
subset_ids = list(range(360))
confidence_scores = torch.rand(len(subset_ids)).tolist()
submission_df = pd.DataFrame({
"id": subset_ids,
"score": confidence_scores
})
submission_df.to_csv("example_submission.csv", index=None)