PyTorch迁移学习Kaggle种子分类：如何保存预测结果为CSV用于提交？

Got it, let's walk through how to turn your trained ResNet50 model's test predictions into a Kaggle-ready CSV file. I'll break this down into straightforward steps that integrate seamlessly with your existing code:

Step 1: Add Prediction Logic for the Test Set #

First, switch your model to evaluation mode (this disables training-specific behaviors like dropout) and run inference on the test set. We'll also disable gradient calculation to save memory and speed things up:

import torch
import pandas as pd

# Switch model to evaluation mode
model.eval()

# Store predictions and corresponding sample IDs (e.g., image filenames/Kaggle test IDs)
predictions = []
sample_ids = []

# Disable gradient computation - no need for backprop during inference
with torch.no_grad():
    # Assuming your test_loader returns (input_tensor, sample_id)
    for inputs, ids in loaders['test']:
        if use_gpu:
            inputs = inputs.cuda()
        
        # Get model outputs
        outputs = model(inputs)
        # Grab the class index with the highest probability
        _, preds = torch.max(outputs, 1)
        
        # Move predictions to CPU and convert to numpy for easy handling
        predictions.extend(preds.cpu().numpy())
        sample_ids.extend(ids)

Quick note: If your current test_loader only returns input tensors, you'll need to update your test Dataset class to return sample IDs too. For example:

from torch.utils.data import Dataset
from PIL import Image
import os
import glob

class SeedTestDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        self.img_paths = sorted(glob.glob(os.path.join(img_dir, "*.jpg")))  # Adjust file type as needed
        self.transform = transform
    
    def __len__(self):
        return len(self.img_paths)
    
    def __getitem__(self, idx):
        img_path = self.img_paths[idx]
        img = Image.open(img_path).convert('RGB')
        if self.transform:
            img = self.transform(img)
        # Extract sample ID from filename (e.g., "test_123.jpg" becomes "test_123")
        sample_id = os.path.basename(img_path).split('.')[0]
        return img, sample_id

Step 2: Format Data for Kaggle Submission #

Kaggle submissions almost always require two columns: one for the sample ID (matching the test set) and one for your predicted label. Convert your collected data into a pandas DataFrame:

# Convert prediction indices to actual class labels using your `classes` list
predicted_labels = [classes[pred_idx] for pred_idx in predictions]

# Create DataFrame with columns matching Kaggle's submission requirements
# Double-check column names (e.g., Kaggle might ask for "ImageId" instead of "Id")
submission_df = pd.DataFrame({
    'Id': sample_ids,
    'Label': predicted_labels
})

Step 3: Save the CSV #

Finally, save the DataFrame to a CSV file. Make sure to exclude the index column—Kaggle doesn't want that!

submission_df.to_csv('seed_classification_submission.csv', index=False)

Full Integrated Code #

Here's how this fits right after your existing training code:

# Your original training code
use_gpu = torch.cuda.is_available()
model = models.resnet50(pretrained=True)
for param in model.parameters():
    param.requires_grad = False
num_ftrs = model.fc.in_features
model.fc = torch.nn.Linear(num_ftrs, len(classes))
if use_gpu:
    model = model.cuda()
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
loaders = {'train':train_loader, 'valid':valid_loader, 'test': test_loader}
model = train_model(loaders, model, criterion, optimizer, exp_lr_scheduler, num_epochs=50)

# New prediction and submission code
model.eval()
predictions = []
sample_ids = []

with torch.no_grad():
    for inputs, ids in loaders['test']:
        if use_gpu:
            inputs = inputs.cuda()
        outputs = model(inputs)
        _, preds = torch.max(outputs, 1)
        predictions.extend(preds.cpu().numpy())
        sample_ids.extend(ids)

predicted_labels = [classes[pred] for pred in predictions]
submission_df = pd.DataFrame({
    'Id': sample_ids,
    'Label': predicted_labels
})
submission_df.to_csv('seed_classification_submission.csv', index=False)

Key Checks Before Submission #

• Match Kaggle's Format: Double-check that your column names and label values exactly match what Kaggle specifies (e.g., if they want numeric labels instead of class names, skip converting to classes and just use predictions directly).
• Sample ID Accuracy: Ensure your sample_ids are identical to the test set IDs provided by Kaggle—mismatches will cause your submission to fail.
• No Extra Index: Always use index=False when saving the CSV to avoid adding an unwanted index column.

内容的提问来源于stack exchange，提问作者Ryan