Product · Sum

import os
import csv
from pathlib import Path
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np
from PIL import Image
from tqdm import tqdm

import random

seed = 42

random.seed(seed)                  # Python built-in random
np.random.seed(seed)               # NumPy
torch.manual_seed(seed)            # PyTorch (CPU)
torch.cuda.manual_seed(seed)       # PyTorch (single GPU)
torch.cuda.manual_seed_all(seed)   # PyTorch (all GPUs)

# Ensures deterministic behavior
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

#-------------Read the training set; the training set path is already set, so the section below does not need to be changed------------------#
#-----Read the training set, the address of the training set has been set, and the following section does not need to be modified-------#
train_path = "/bohr/train-rppd/v1"

# Read the data.
def load_train_data(data_dir='./train/'):
    label_path = os.path.join(data_dir, 'train_labels.csv')
    image_dir = os.path.join(data_dir, 'train_images')

    data = []
    with open(label_path, 'r', newline='', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        for row in reader:
            sample_id = row['id']
            data.append({
                'id': sample_id,
                'image_path': os.path.join(image_dir, f'{sample_id}.png'),
                'sum': float(row['sum']),
                'product': float(row['product']),
            })

    print(f'Successfully loaded training records: {len(data)}')
    return data


class MNISTBaselineDataset(Dataset):

    def __init__(self, records):
        self.records = records

    def __len__(self):
        return len(self.records)

    @staticmethod
    def load_image(image_path):
        image = Image.open(image_path).convert('L')
        image = np.array(image, dtype=np.float32) / 255.0
        image = (image - 0.1307) / 0.3081
        return image

    def __getitem__(self, idx):
        record = self.records[idx]
        image = self.load_image(record['image_path'])
        image = torch.tensor(image, dtype=torch.float32).unsqueeze(0)
        target = torch.tensor([record['sum'], record['product']], dtype=torch.float32)
        return image, target


def create_train_loader(batch_size=64, num_workers=1, data_dir='./train/'):
    records = load_train_data(data_dir)
    dataset = MNISTBaselineDataset(records)
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
    )


class SimpleCNN(nn.Module):

    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1)),
        )
        self.regressor = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 2),
        )

    def forward(self, x):
        x = self.features(x)
        return self.regressor(x)


# Train the model
def train(model, train_loader, epochs, device='cpu'):
    model.to(device)
    criterion = nn.L1Loss()
    optimizer = optim.SGD(model.parameters(), lr=0.01)

    for epoch in range(1, epochs + 1):
        model.train()
        total_loss = 0.0

        for images, targets in tqdm(train_loader, desc=f'Epoch {epoch}/{epochs}', leave=False):
            images = images.to(device)
            targets = targets.to(device)

            optimizer.zero_grad()
            preds = model(images)
            loss = criterion(preds, targets)
            loss.backward()
            optimizer.step()
            total_loss += loss.item() * images.size(0)

        avg_loss = total_loss / len(train_loader.dataset)
        print(f'Epoch {epoch}/{epochs} - Loss: {avg_loss:.4f}')


def set_random_seed(seed=42):
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)


device = 'cuda' if torch.cuda.is_available() else 'cpu'
seed = 42
batch_size = 64
epochs = 20

set_random_seed(seed)
train_loader = create_train_loader(batch_size=batch_size, data_dir=train_path)
model = SimpleCNN()
train(model, train_loader, epochs=epochs, device=str(device))

#-------------Read the test set---------------#“DATA_PATH” is the environment variable of the encrypted test set; in the way shown below, the test set can be accessed during system scoring after submission, but contestants cannot download it directly
#----Read the testing set, “DATA_PATH” is an environment variable for the encrypted test set. After submission, you can access the test set for system scoring in the following manner, but the contestant cannot download it directly.-----#
if os.environ.get('DATA_PATH'):
    test_path = os.environ.get("DATA_PATH") + "/"
else:
    test_path = "./test/"
    print("When the baseline runs, the test set cannot be read, so this error message appears; this is normal.")
    print("When baseline is running, this error message will appear because the test set cannot be read, which is a normal phenomenon.")
    #When the baseline runs, the test set cannot be read, so this error message appears; this is normal.
    #When baseline is running, this error message will appear because the test set cannot be read, which is a normal phenomenon.

# Read the test data
class MNISTTestDataset(Dataset):
    def __init__(self, image_dir):
        self.image_paths = sorted(str(path) for path in Path(image_dir).glob('*.png'))

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        image_path = self.image_paths[idx]
        sample_id = Path(image_path).stem
        image = MNISTBaselineDataset.load_image(image_path)
        image = torch.tensor(image, dtype=torch.float32).unsqueeze(0)
        return image, sample_id


# Here the trained model directly regresses sum and product.
def predict_and_save(model, image_dir, output_file, device='cpu', batch_size=64, num_workers=1):
    dataset = MNISTTestDataset(image_dir)
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
    )

    rows = []
    model.eval()
    with torch.no_grad():
        for images, sample_ids in tqdm(loader, desc=f'Predicting {os.path.basename(image_dir)}', leave=False):
            images = images.to(device)
            preds = model(images).cpu()
            sum_preds = preds[:, 0].round().clamp(0, 36).to(torch.int64).numpy()
            product_preds = preds[:, 1].round().clamp(0, 6561).to(torch.int64).numpy()

            for sample_id, sum_pred, product_pred in zip(sample_ids, sum_preds, product_preds):
                rows.append({
                    'id': sample_id,
                    'sum': int(sum_pred),
                    'product': int(product_pred),
                })

    with open(output_file, 'w', newline='', encoding='utf-8') as f:
        writer = csv.DictWriter(f, fieldnames=['id', 'sum', 'product'])
        writer.writeheader()
        writer.writerows(rows)

    print(f'Saved predictions to {output_file}')


val_dir = os.path.join(test_path, 'val')
test_dir = os.path.join(test_path, 'test')

for required_dir in [val_dir, test_dir]:
    if not os.path.isdir(required_dir):
        raise FileNotFoundError(f'Missing test directory: {required_dir}')

predict_and_save(model, val_dir, 'submission_val.csv', device=str(device), batch_size=batch_size)
predict_and_save(model, test_dir, 'submission_test.csv', device=str(device), batch_size=batch_size)

import zipfile

# Define the files to pack and the name of the archive
files_to_zip = ['submission_val.csv', 'submission_test.csv']
zip_filename = 'submission.zip'

# Create a zip file
with zipfile.ZipFile(zip_filename, 'w') as zipf:
    for file in files_to_zip:
        # Add the file to the zip file
        zipf.write(file, os.path.basename(file))

print(f'{zip_filename} created successfully!')