Import Dependencies

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader, random_split

Training Configuration Setup

Customize these hyperparameters to tune model performance and ensure reproducibility.

device = "cuda" if torch.cuda.is_available() else "cpu"
training_config = {
    "random_seed": 999123,  # Custom random seed for consistent results
    "use_full_features": True,  # Toggle to utilize all input features
    "val_split": 0.15,  # Fraction of training data allocated for validation
    "total_epochs": 4000,  # Number of complete training cycles
    "batch_size": 128,
    "learning_rate": 5e-6,
    "patience": 500,  # Early stopping threshold for stalled validation loss
    "checkpoint_path": "./checkpoints/best_model.ckpt"  # Path to save the optimal model
}

Define the Regression Neural Network

This feedforward network maps input features to continuous output values using fully connected layers and ReLU activations.

class RegressionModel(nn.Module):
    def __init__(self, input_features):
        super().__init__()
        # Stack of fully connected layers with non-linear activations
        self.fc_stack = nn.Sequential(
            nn.Linear(input_features, 32),
            nn.ReLU(),
            nn.Linear(32, 16),
            nn.ReLU(),
            nn.Linear(16, 1)
        )
    
    def forward(self, x):
        output = self.fc_stack(x)
        return output.squeeze(dim=1)  # Reduce tensor dimension from (batch_size, 1) to (batch_size)

Model Training Loop

This function handles training, validation, early stopping, and model checkpointing.

def train_model(train_dataloader, val_dataloader, model, config, compute_device):
    loss_function = nn.MSELoss(reduction="mean")  # Mean Squared Loss for regression tasks
    # Stochastic Gradient Descent optimizer with momentum
    optimizer = torch.optim.SGD(model.parameters(), lr=config["learning_rate"], momentum=0.95)
    
    total_epochs = config["total_epochs"]
    best_val_loss = float("inf")
    step_count = 0
    stop_tracker = 0

    for epoch in range(total_epochs):
        model.train()  # Enable training mode (activates dropout/batch norm updates)
        training_losses = []

        for batch_x, batch_y in train_dataloader:
            optimizer.zero_grad()  # Reset gradients to prevent accumulation
            batch_x, batch_y = batch_x.to(compute_device), batch_y.to(compute_device)
            
            pred = model(batch_x)
            loss = loss_function(pred, batch_y)
            
            loss.backward()  # Compute gradients via backpropagation
            optimizer.step()  # Update model parameters
            
            step_count += 1
            training_losses.append(loss.detach().item())
            
            # Update progress bar with current training status
            train_pbar.set_description(f"Epoch [{epoch+1}/{total_epochs}]")
            train_pbar.set_postfix({"Loss": loss.detach().item()})
        
        avg_train_loss = sum(training_losses) / len(training_losses)
        writer.add_scalar("Loss/Train", avg_train_loss, step_count)

        model.eval()  # Switch to evaluation mode (disables training-specific layers)
        validation_losses = []
        with torch.no_grad():  # Disable gradient computation to save memory
            for batch_x, batch_y in val_dataloader:
                batch_x, batch_y = batch_x.to(compute_device), batch_y.to(compute_device)
                pred = model(batch_x)
                val_loss = loss_function(pred, batch_y)
                validation_losses.append(val_loss.item())
        
        avg_val_loss = sum(validation_losses) / len(validation_losses)
        print(f"Epoch [{epoch+1}/{total_epochs}]: Avg Train Loss: {avg_train_loss:.4f}, Avg Val Loss: {avg_val_loss:.4f}")
        writer.add_scalar("Loss/Validation", avg_val_loss, step_count)

        # Save the model if it achieves the lowest validation loss so far
        if avg_val_loss < best_val_loss:
            best_val_loss = avg_val_loss
            torch.save(model.state_dict(), config["checkpoint_path"])
            print(f"Best model saved with validation loss: {best_val_loss:.3f}")
            stop_tracker = 0
        else:
            stop_tracker += 1
        
        # Early stopping if no improvement for specified number of epochs
        if stop_tracker >= config["patience"]:
            print("\nTraining halted early due to no validation loss improvement.")
            return

Initialize and Launch Training

# Determine input feature count from training dataset
input_feature_count = x_train.shape[1]
model = RegressionModel(input_feature_count).to(device)

# Begin training process
train_model(train_loader, val_loader, model, training_config, device)

Inference Function for Test Data

def run_inference(test_dataloader, model, compute_device):
    model.eval()
    predictions = []
    
    for batch in tqdm(test_dataloader):
        batch = batch.to(compute_device)
        with torch.no_grad():
            batch_preds = model(batch)
            predictions.append(batch_preds.detach().cpu())
    
    return torch.cat(predictions, dim=0).numpy()

Generate and Export Test Predictions

def export_predictions(predictions, output_file_path):
    """Save prediction results to a CSV file in the required format"""
    import csv
    with open(output_file_path, "w", newline="") as file:
        writer = csv.writer(file)
        writer.writerow(["id", "predicted_value"])
        for index, pred in enumerate(predictions):
            writer.writerow([index, pred])

# Load the pre-trained best model
inference_model = RegressionModel(input_feature_count).to(device)
inference_model.load_state_dict(torch.load(training_config["checkpoint_path"]))

# Generate predictions on test dataset
test_results = run_inference(test_loader, inference_model, device)

# Export results to CSV
export_predictions(test_results, "test_predictions.csv")