train.py

"""
PyTorch Training Template
=========================

A comprehensive training template that supports all PyTorch algorithms, optimizers,
loss functions, and activation functions. Configure everything via global variables.

Available Algorithms:
    - SimpleNet, MLP: Feedforward neural networks
    - Linear: Simple linear model
    - CNN, ConvNet: Convolutional Neural Networks
    - RNN: Recurrent Neural Networks
    - LSTM: Long Short-Term Memory
    - GRU: Gated Recurrent Unit
    - Transformer: Transformer models

Available Optimizers (from torch.optim):
    - SGD, Adam, AdamW, RMSprop, Adagrad, Adadelta, Adamax, ASGD, LBFGS,
      Rprop, RAdam, NAdam, SparseAdam

Available Loss Functions (from torch.nn):
    Regression: MSELoss, L1Loss, SmoothL1Loss, HuberLoss, PoissonNLLLoss,
                GaussianNLLLoss, KLDivLoss
    Classification: CrossEntropyLoss, BCELoss, BCEWithLogitsLoss, NLLLoss,
                     MultiLabelMarginLoss, MultiLabelSoftMarginLoss,
                     MultiMarginLoss, SoftMarginLoss, MarginRankingLoss,
                     TripletMarginLoss, HingeEmbeddingLoss, CTCLoss
    Other: CosineEmbeddingLoss

Available Activations (from torch.nn):
    - ReLU, ReLU6, LeakyReLU, PReLU, RReLU, GELU, Sigmoid, Tanh, Hardtanh,
      Hardswish, ELU, CELU, SELU, GLU, SiLU, Mish, Softplus, Softshrink,
      Hardshrink, Softsign, Tanhshrink, Threshold, Hardsigmoid, LogSigmoid,
      Softmin, Softmax, LogSoftmax, Identity

Reference: https://docs.pytorch.org/docs/stable/nn.html
"""

import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as Data
import matplotlib.pyplot as plt
import numpy as np
import random
import time
import os
import sys
import logging
from logging.handlers import RotatingFileHandler
from pathlib import Path

# ============================================================================
# GLOBAL CONFIGURATION VARIABLES
# ============================================================================
# Model Configuration
MODEL_VERSION = 1
MODEL_NAME = None            # Model name (used in filenames). If None, uses model_v{MODEL_VERSION}
ALGORITHM = "SimpleNet"     # Model architecture: SimpleNet, Linear, MLP, CNN, RNN, LSTM, GRU, Transformer, etc.
INPUT_SIZE = 1              # Input feature size (or input channels for CNN)
HIDDEN_SIZES = [10]         # List of hidden layer sizes
OUTPUT_SIZE = 1             # Output size
ACTIVATION = "ReLU"         # Activation function (see get_activation() for all options)

# CNN Configuration (if ALGORITHM == "CNN")
CNN_INPUT_CHANNELS = 1      # Input channels (e.g., 1 for grayscale, 3 for RGB)
CNN_OUTPUT_CHANNELS = [32, 64]  # List of output channels for each conv layer
CNN_KERNEL_SIZES = [3, 3]  # Kernel sizes for each conv layer
CNN_STRIDES = [1, 1]        # Strides for each conv layer
CNN_PADDING = [1, 1]        # Padding for each conv layer
CNN_POOL_KERNEL = 2         # Pooling kernel size
CNN_DROPOUT = 0.0           # Dropout rate for CNN

# RNN/LSTM/GRU Configuration (if ALGORITHM in ["RNN", "LSTM", "GRU"])
RNN_HIDDEN_SIZE = 64        # Hidden size for RNN/LSTM/GRU
RNN_NUM_LAYERS = 1          # Number of RNN layers
RNN_BIDIRECTIONAL = False   # Whether to use bidirectional RNN
RNN_DROPOUT = 0.0           # Dropout rate for RNN layers
RNN_SEQUENCE_LENGTH = 10    # Sequence length (if not provided in data)

# Transformer Configuration (if ALGORITHM == "Transformer")
TRANSFORMER_D_MODEL = 512   # Model dimension
TRANSFORMER_NHEAD = 8       # Number of attention heads
TRANSFORMER_NUM_LAYERS = 6  # Number of transformer layers
TRANSFORMER_DIM_FEEDFORWARD = 2048  # Feedforward dimension
TRANSFORMER_DROPOUT = 0.1   # Dropout rate
TRANSFORMER_MAX_SEQ_LEN = 100  # Maximum sequence length

# Training Configuration
BATCH_SIZE = 100
EPOCHS = 100
LEARNING_RATE = 0.01
SEED = 42
OPTIMIZER = "Adam"          # Optimizer: Adam, SGD, AdamW, RMSprop, Adagrad, Adadelta, etc.
LOSS_FUNCTION = "MSELoss"   # Loss: MSELoss, L1Loss, SmoothL1Loss, CrossEntropyLoss, BCELoss, etc.
WEIGHT_DECAY = 0.0          # L2 regularization
MOMENTUM = 0.9              # For SGD optimizer
NESTEROV = False            # For SGD with Nesterov momentum
BETAS = (0.9, 0.999)        # For Adam/AdamW optimizers
EPS = 1e-8                  # Epsilon for optimizers
AMSGRAD = False             # For Adam optimizer
ALPHA = 0.99                # For RMSprop optimizer
CENTERED = False            # For RMSprop optimizer (centered parameter)
RHO = 0.9                   # For Adadelta optimizer
LR_DECAY = 0               # For Adagrad optimizer

# Metrics Configuration
METRICS = ["MAE", "R2"]             # Metrics to track: MAE, R2, Accuracy, etc.
METRICS_THRESHOLD_VALUE = 0.9       # Threshold value for metrics
METRICS_THRESHOLD_TYPE = "max"      # "max" or "min" - whether higher/lower is better

# Device Configuration
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Logging Configuration
LOG_INTERVAL = 10              # Log every N epochs
SAVE_INTERVAL = 10             # Save model every N epochs
SAVE_PATH = "models"           # Path to save models
LOG_PATH = "logs"              # Path to save logs
# LOG_FILE will be set based on MODEL_NAME or MODEL_VERSION
LOG_FILE = None  # Will be set to train_{MODEL_NAME}.log or train_{MODEL_VERSION}.log
LOG_LEVEL = "INFO"             # DEBUG, INFO, WARNING, ERROR
LOG_FORMAT = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"  # Log format
LOG_DATE_FORMAT = "%Y-%m-%d %H:%M:%S"  # Log date format
LOG_FILE_MAX_BYTES = 1024 * 1024 * 10  # 10MB
LOG_FILE_BACKUP_COUNT = 10     # Number of backup files to keep

# Data Configuration
TRAIN_TEST_SPLIT = 0.8          # 80% train, 20% test
SHUFFLE_DATA = True             # Whether to shuffle data
NORMALIZE_DATA = False          # Whether to normalize input data

# Early Stopping Configuration
EARLY_STOPPING = False             # Whether to use early stopping
EARLY_STOPPING_PATIENCE = 10       # Number of epochs to wait before stopping
EARLY_STOPPING_METRIC = "loss"     # Metric to check for early stopping
EARLY_STOPPING_MIN_DELTA = 0.001   # Minimum change in metric to consider improvement

# Model Saving/Loading Configuration
SAVE_FORMAT = "pth"                # File format: "pth" (PyTorch pickle), "pt" (same as pth)
SAVE_WEIGHTS_ONLY = False          # If True, save only state_dict (safer, but requires model architecture to recreate)
LOAD_WEIGHTS_ONLY = False          # If True, use weights_only=True when loading (PyTorch 1.13.0+, safer for untrusted files)
PICKLE_PROTOCOL = None             # Pickle protocol version (None = default, 2-5 supported)

# ============================================================================
# SETUP FUNCTIONS
# ============================================================================

def setup_seed(seed):
    """Set random seeds for reproducibility."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

def setup_directories():
    """Create necessary directories if they don't exist."""
    Path(SAVE_PATH).mkdir(parents=True, exist_ok=True)
    Path(LOG_PATH).mkdir(parents=True, exist_ok=True)

def setup_logging():
    """Configure logging to file and console."""
    # Set LOG_FILE if not already set
    global LOG_FILE
    if LOG_FILE is None:
        if MODEL_NAME:
            LOG_FILE = f"train_{MODEL_NAME}.log"
        else:
            LOG_FILE = f"train_{MODEL_VERSION}.log"
    
    # Create logger
    logger = logging.getLogger()
    logger.setLevel(getattr(logging, LOG_LEVEL))
    
    # Clear existing handlers
    logger.handlers = []
    
    # File handler with rotation
    log_file_path = os.path.join(LOG_PATH, LOG_FILE)
    file_handler = RotatingFileHandler(
        log_file_path,
        maxBytes=LOG_FILE_MAX_BYTES,
        backupCount=LOG_FILE_BACKUP_COUNT
    )
    file_handler.setLevel(getattr(logging, LOG_LEVEL))
    file_formatter = logging.Formatter(LOG_FORMAT, datefmt=LOG_DATE_FORMAT)
    file_handler.setFormatter(file_formatter)
    logger.addHandler(file_handler)
    
    # Console handler
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(getattr(logging, LOG_LEVEL))
    console_formatter = logging.Formatter(LOG_FORMAT, datefmt=LOG_DATE_FORMAT)
    console_handler.setFormatter(console_formatter)
    logger.addHandler(console_handler)
    
    return logger

# ============================================================================
# DATA PROCESSING FUNCTIONS
# ============================================================================

def convert_to_tensor(data):
    """Convert numpy array to torch tensor if needed."""
    if isinstance(data, np.ndarray):
        return torch.from_numpy(data).float()
    elif isinstance(data, torch.Tensor):
        return data.float()
    else:
        raise TypeError(f"Data must be numpy array or torch tensor, got {type(data)}")

def normalize_data(data, mean=None, std=None):
    """Normalize data to have mean 0 and std 1."""
    if mean is None:
        mean = data.mean()
    if std is None:
        std = data.std()
    normalized = (data - mean) / (std + 1e-8)  # Add small epsilon to avoid division by zero
    return normalized, mean, std

def prepare_data(X, y, train_split=None):
    """
    Prepare data for training.
    
    Args:
        X: Input features (numpy array or tensor)
        y: Target values (numpy array or tensor)
        train_split: Train/test split ratio (uses global TRAIN_TEST_SPLIT if None)
    
    Returns:
        train_loader, test_loader, data_stats
    """
    logger = logging.getLogger()
    logger.info("Preparing data...")
    
    # Convert to tensors
    X = convert_to_tensor(X)
    y = convert_to_tensor(y)
    
    # Normalize if requested
    data_stats = {}
    if NORMALIZE_DATA:
        X, mean, std = normalize_data(X)
        data_stats['X_mean'] = mean.item()
        data_stats['X_std'] = std.item()
        logger.info(f"Normalized X: mean={mean.item():.4f}, std={std.item():.4f}")
    
    # Ensure proper shapes
    if len(X.shape) == 1:
        X = X.unsqueeze(1)
    if len(y.shape) == 1:
        y = y.unsqueeze(1)
    
    # Train/test split
    split_ratio = train_split if train_split is not None else TRAIN_TEST_SPLIT
    n_train = int(len(X) * split_ratio)
    
    if SHUFFLE_DATA:
        indices = torch.randperm(len(X))
        X = X[indices]
        y = y[indices]
    
    X_train, X_test = X[:n_train], X[n_train:]
    y_train, y_test = y[:n_train], y[n_train:]
    
    # Create data loaders
    train_dataset = Data.TensorDataset(X_train, y_train)
    test_dataset = Data.TensorDataset(X_test, y_test)
    
    train_loader = Data.DataLoader(
        train_dataset,
        batch_size=BATCH_SIZE,
        shuffle=True
    )
    test_loader = Data.DataLoader(
        test_dataset,
        batch_size=BATCH_SIZE,
        shuffle=False
    )
    
    logger.info(f"Train samples: {len(X_train)}, Test samples: {len(X_test)}")
    logger.info(f"Input shape: {X_train.shape[1:]}, Output shape: {y_train.shape[1:]}")
    
    return train_loader, test_loader, data_stats

# ============================================================================
# MODEL DEFINITION
# ============================================================================

def get_activation(activation_name, **kwargs):
    """
    Get activation function by name.
    Supports all PyTorch activation functions from torch.nn.
    Reference: https://docs.pytorch.org/docs/stable/nn.html#non-linear-activations-weighted-sum-nonlinearity
    """
    # Handle activations that require specific parameters
    if activation_name == "Threshold":
        threshold = kwargs.get('threshold', 0.0)
        value = kwargs.get('value', 0.0)
        return nn.Threshold(threshold, value)
    elif activation_name == "PReLU":
        num_parameters = kwargs.get('num_parameters', 1)
        init = kwargs.get('init', 0.25)
        return nn.PReLU(num_parameters=num_parameters, init=init)
    elif activation_name == "RReLU":
        lower = kwargs.get('lower', 1.0/8)
        upper = kwargs.get('upper', 1.0/3)
        return nn.RReLU(lower=lower, upper=upper)
    elif activation_name == "LeakyReLU":
        negative_slope = kwargs.get('negative_slope', 0.01)
        return nn.LeakyReLU(negative_slope=negative_slope)
    elif activation_name == "Hardtanh":
        min_val = kwargs.get('min_val', -1.0)
        max_val = kwargs.get('max_val', 1.0)
        return nn.Hardtanh(min_val=min_val, max_val=max_val)
    elif activation_name == "ELU":
        alpha = kwargs.get('alpha', 1.0)
        return nn.ELU(alpha=alpha)
    elif activation_name == "CELU":
        alpha = kwargs.get('alpha', 1.0)
        return nn.CELU(alpha=alpha)
    elif activation_name == "Softplus":
        beta = kwargs.get('beta', 1)
        threshold = kwargs.get('threshold', 20)
        return nn.Softplus(beta=beta, threshold=threshold)
    elif activation_name == "Softshrink":
        lambd = kwargs.get('lambd', 0.5)
        return nn.Softshrink(lambd=lambd)
    elif activation_name == "Hardshrink":
        lambd = kwargs.get('lambd', 0.5)
        return nn.Hardshrink(lambd=lambd)
    elif activation_name == "Softmax":
        dim = kwargs.get('dim', None)
        if dim is None:
            return nn.Softmax(dim=1)  # Default dimension
        return nn.Softmax(dim=dim)
    elif activation_name == "LogSoftmax":
        dim = kwargs.get('dim', None)
        if dim is None:
            return nn.LogSoftmax(dim=1)  # Default dimension
        return nn.LogSoftmax(dim=dim)
    elif activation_name == "Softmin":
        dim = kwargs.get('dim', None)
        if dim is None:
            return nn.Softmin(dim=1)  # Default dimension
        return nn.Softmin(dim=dim)
    
    # Standard activations that don't require special parameters
    activations = {
        "ReLU": nn.ReLU(),
        "ReLU6": nn.ReLU6(),
        "GELU": nn.GELU(),
        "Sigmoid": nn.Sigmoid(),
        "Tanh": nn.Tanh(),
        "Hardswish": nn.Hardswish(),
        "SELU": nn.SELU(),
        "GLU": nn.GLU(),
        "SiLU": nn.SiLU(),
        "Mish": nn.Mish(),
        "Softsign": nn.Softsign(),
        "Tanhshrink": nn.Tanhshrink(),
        "Hardsigmoid": nn.Hardsigmoid(),
        "LogSigmoid": nn.LogSigmoid(),
        # Identity (no activation)
        "None": nn.Identity(),
        None: nn.Identity(),
        "Identity": nn.Identity()
    }
    
    return activations.get(activation_name, nn.ReLU())

def create_model():
    """
    Create model based on ALGORITHM global variable.
    Supports: SimpleNet, Linear, MLP, CNN, RNN, LSTM, GRU, Transformer, ResNet, etc.
    Reference: https://docs.pytorch.org/docs/stable/nn.html
    Returns model instance.
    """
    logger = logging.getLogger()
    logger.info(f"Creating model: {ALGORITHM}")
    
    if ALGORITHM == "SimpleNet" or ALGORITHM == "MLP":
        # Simple feedforward network / Multi-Layer Perceptron
        layers = []
        input_dim = INPUT_SIZE
        
        # Build hidden layers
        for hidden_size in HIDDEN_SIZES:
            layers.append(nn.Linear(input_dim, hidden_size))
            layers.append(get_activation(ACTIVATION))
            input_dim = hidden_size
        
        # Output layer
        layers.append(nn.Linear(input_dim, OUTPUT_SIZE))
        model = nn.Sequential(*layers)
        
    elif ALGORITHM == "Linear":
        # Simple linear model
        model = nn.Linear(INPUT_SIZE, OUTPUT_SIZE)
        
    elif ALGORITHM == "CNN" or ALGORITHM == "ConvNet":
        # Convolutional Neural Network
        layers = []
        in_channels = CNN_INPUT_CHANNELS
        
        # Build convolutional layers
        for i, (out_channels, kernel_size, stride, padding) in enumerate(
            zip(CNN_OUTPUT_CHANNELS, CNN_KERNEL_SIZES, CNN_STRIDES, CNN_PADDING)
        ):
            layers.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding))
            layers.append(get_activation(ACTIVATION))
            layers.append(nn.MaxPool2d(CNN_POOL_KERNEL))
            if CNN_DROPOUT > 0:
                layers.append(nn.Dropout(CNN_DROPOUT))
            in_channels = out_channels
        
        # Flatten and add fully connected layers
        layers.append(nn.Flatten())
        # Calculate flattened size (this is approximate - may need adjustment based on input)
        fc_input_size = in_channels * 7 * 7  # Adjust based on your input size
        for hidden_size in HIDDEN_SIZES:
            layers.append(nn.Linear(fc_input_size, hidden_size))
            layers.append(get_activation(ACTIVATION))
            if CNN_DROPOUT > 0:
                layers.append(nn.Dropout(CNN_DROPOUT))
            fc_input_size = hidden_size
        layers.append(nn.Linear(fc_input_size, OUTPUT_SIZE))
        model = nn.Sequential(*layers)
        
    elif ALGORITHM == "RNN":
        # Recurrent Neural Network
        class RNNModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.rnn = nn.RNN(
                    INPUT_SIZE, RNN_HIDDEN_SIZE, RNN_NUM_LAYERS,
                    batch_first=True, bidirectional=RNN_BIDIRECTIONAL,
                    dropout=RNN_DROPOUT if RNN_NUM_LAYERS > 1 else 0
                )
                self.fc = nn.Linear(
                    RNN_HIDDEN_SIZE * (2 if RNN_BIDIRECTIONAL else 1),
                    OUTPUT_SIZE
                )
            
            def forward(self, x):
                # x shape: (batch, seq_len, features) or (batch, features)
                if len(x.shape) == 2:
                    x = x.unsqueeze(1)  # Add sequence dimension
                out, _ = self.rnn(x)
                # Take the last output
                out = out[:, -1, :]
                return self.fc(out)
        model = RNNModel()
        
    elif ALGORITHM == "LSTM":
        # Long Short-Term Memory
        class LSTMModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.lstm = nn.LSTM(
                    INPUT_SIZE, RNN_HIDDEN_SIZE, RNN_NUM_LAYERS,
                    batch_first=True, bidirectional=RNN_BIDIRECTIONAL,
                    dropout=RNN_DROPOUT if RNN_NUM_LAYERS > 1 else 0
                )
                self.fc = nn.Linear(
                    RNN_HIDDEN_SIZE * (2 if RNN_BIDIRECTIONAL else 1),
                    OUTPUT_SIZE
                )
            
            def forward(self, x):
                if len(x.shape) == 2:
                    x = x.unsqueeze(1)
                out, _ = self.lstm(x)
                out = out[:, -1, :]
                return self.fc(out)
        model = LSTMModel()
        
    elif ALGORITHM == "GRU":
        # Gated Recurrent Unit
        class GRUModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.gru = nn.GRU(
                    INPUT_SIZE, RNN_HIDDEN_SIZE, RNN_NUM_LAYERS,
                    batch_first=True, bidirectional=RNN_BIDIRECTIONAL,
                    dropout=RNN_DROPOUT if RNN_NUM_LAYERS > 1 else 0
                )
                self.fc = nn.Linear(
                    RNN_HIDDEN_SIZE * (2 if RNN_BIDIRECTIONAL else 1),
                    OUTPUT_SIZE
                )
            
            def forward(self, x):
                if len(x.shape) == 2:
                    x = x.unsqueeze(1)
                out, _ = self.gru(x)
                out = out[:, -1, :]
                return self.fc(out)
        model = GRUModel()
        
    elif ALGORITHM == "Transformer":
        # Transformer model
        class TransformerModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.embedding = nn.Linear(INPUT_SIZE, TRANSFORMER_D_MODEL)
                encoder_layer = nn.TransformerEncoderLayer(
                    d_model=TRANSFORMER_D_MODEL,
                    nhead=TRANSFORMER_NHEAD,
                    dim_feedforward=TRANSFORMER_DIM_FEEDFORWARD,
                    dropout=TRANSFORMER_DROPOUT,
                    batch_first=True
                )
                self.transformer = nn.TransformerEncoder(
                    encoder_layer,
                    num_layers=TRANSFORMER_NUM_LAYERS
                )
                self.fc = nn.Linear(TRANSFORMER_D_MODEL, OUTPUT_SIZE)
            
            def forward(self, x):
                if len(x.shape) == 2:
                    x = x.unsqueeze(1)  # Add sequence dimension
                x = self.embedding(x)
                x = self.transformer(x)
                x = x[:, -1, :]  # Take last sequence element
                return self.fc(x)
        model = TransformerModel()
        
    else:
        # Default: SimpleNet
        logger.warning(f"Unknown algorithm {ALGORITHM}, using SimpleNet")
        layers = []
        input_dim = INPUT_SIZE
        for hidden_size in HIDDEN_SIZES:
            layers.append(nn.Linear(input_dim, hidden_size))
            layers.append(get_activation(ACTIVATION))
            input_dim = hidden_size
        layers.append(nn.Linear(input_dim, OUTPUT_SIZE))
        model = nn.Sequential(*layers)
    
    model = model.to(DEVICE)
    
    # Log model info
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    logger.info(f"Model created with {total_params} total parameters ({trainable_params} trainable)")
    logger.info(f"Model architecture:\n{model}")
    
    return model

# ============================================================================
# OPTIMIZER AND LOSS FUNCTION
# ============================================================================

def get_optimizer(model):
    """
    Get optimizer based on OPTIMIZER global variable.
    Supports all PyTorch optimizers from torch.optim.
    Reference: https://docs.pytorch.org/docs/stable/optim.html
    """
    optimizers = {
        # Stochastic Gradient Descent
        "SGD": optim.SGD(
            model.parameters(), lr=LEARNING_RATE, momentum=MOMENTUM,
            weight_decay=WEIGHT_DECAY, nesterov=NESTEROV
        ),
        # Adam optimizers
        "Adam": optim.Adam(
            model.parameters(), lr=LEARNING_RATE, betas=BETAS,
            eps=EPS, weight_decay=WEIGHT_DECAY, amsgrad=AMSGRAD
        ),
        "AdamW": optim.AdamW(
            model.parameters(), lr=LEARNING_RATE, betas=BETAS,
            eps=EPS, weight_decay=WEIGHT_DECAY, amsgrad=AMSGRAD
        ),
        # RMSprop
        "RMSprop": optim.RMSprop(
            model.parameters(), lr=LEARNING_RATE, alpha=ALPHA,
            eps=EPS, weight_decay=WEIGHT_DECAY, momentum=MOMENTUM,
            centered=CENTERED
        ),
        # Adagrad
        "Adagrad": optim.Adagrad(
            model.parameters(), lr=LEARNING_RATE, lr_decay=LR_DECAY,
            weight_decay=WEIGHT_DECAY, eps=EPS
        ),
        # Adadelta
        "Adadelta": optim.Adadelta(
            model.parameters(), lr=LEARNING_RATE, rho=RHO,
            eps=EPS, weight_decay=WEIGHT_DECAY
        ),
        # Adamax
        "Adamax": optim.Adamax(
            model.parameters(), lr=LEARNING_RATE, betas=BETAS,
            eps=EPS, weight_decay=WEIGHT_DECAY
        ),
        # ASGD (Averaged SGD)
        "ASGD": optim.ASGD(
            model.parameters(), lr=LEARNING_RATE, lambd=0.0001,
            alpha=0.75, t0=1000000.0, weight_decay=WEIGHT_DECAY
        ),
        # LBFGS
        "LBFGS": optim.LBFGS(
            model.parameters(), lr=LEARNING_RATE, max_iter=20,
            max_eval=None, tolerance_grad=1e-07, tolerance_change=1e-09,
            history_size=100, line_search_fn=None
        ),
        # Rprop
        "Rprop": optim.Rprop(
            model.parameters(), lr=LEARNING_RATE, etas=(0.5, 1.2),
            step_sizes=(1e-06, 50)
        ),
        # RAdam
        "RAdam": optim.RAdam(
            model.parameters(), lr=LEARNING_RATE, betas=BETAS,
            eps=EPS, weight_decay=WEIGHT_DECAY
        ),
        # NAdam
        "NAdam": optim.NAdam(
            model.parameters(), lr=LEARNING_RATE, betas=BETAS,
            eps=EPS, weight_decay=WEIGHT_DECAY, momentum_decay=0.004
        ),
        # SparseAdam
        "SparseAdam": optim.SparseAdam(
            model.parameters(), lr=LEARNING_RATE, betas=BETAS, eps=EPS
        )
    }
    
    optimizer = optimizers.get(OPTIMIZER)
    if optimizer is None:
        logger = logging.getLogger()
        logger.warning(f"Unknown optimizer {OPTIMIZER}, using Adam")
        optimizer = optim.Adam(
            model.parameters(), lr=LEARNING_RATE, betas=BETAS,
            eps=EPS, weight_decay=WEIGHT_DECAY
        )
    
    return optimizer

def get_loss_function():
    """
    Get loss function based on LOSS_FUNCTION global variable.
    Supports all PyTorch loss functions from torch.nn.
    Reference: https://docs.pytorch.org/docs/stable/nn.html#loss-functions
    """
    losses = {
        # Regression losses
        "MSELoss": nn.MSELoss(),
        "L1Loss": nn.L1Loss(),
        "SmoothL1Loss": nn.SmoothL1Loss(),
        "HuberLoss": nn.HuberLoss(),
        "PoissonNLLLoss": nn.PoissonNLLLoss(),
        "GaussianNLLLoss": nn.GaussianNLLLoss(),
        "KLDivLoss": nn.KLDivLoss(),
        
        # Classification losses
        "CrossEntropyLoss": nn.CrossEntropyLoss(),
        "BCELoss": nn.BCELoss(),
        "BCEWithLogitsLoss": nn.BCEWithLogitsLoss(),
        "NLLLoss": nn.NLLLoss(),
        "MultiLabelMarginLoss": nn.MultiLabelMarginLoss(),
        "MultiLabelSoftMarginLoss": nn.MultiLabelSoftMarginLoss(),
        "MultiMarginLoss": nn.MultiMarginLoss(),
        "SoftMarginLoss": nn.SoftMarginLoss(),
        "MarginRankingLoss": nn.MarginRankingLoss(),
        "TripletMarginLoss": nn.TripletMarginLoss(),
        "TripletMarginWithDistanceLoss": nn.TripletMarginWithDistanceLoss(),
        "HingeEmbeddingLoss": nn.HingeEmbeddingLoss(),
        "CTCLoss": nn.CTCLoss(),
        
        # Other losses
        "CosineEmbeddingLoss": nn.CosineEmbeddingLoss(),
        "LabelSmoothingCrossEntropy": nn.CrossEntropyLoss(label_smoothing=0.1)
    }
    
    loss_fn = losses.get(LOSS_FUNCTION)
    if loss_fn is None:
        logger = logging.getLogger()
        logger.warning(f"Unknown loss function {LOSS_FUNCTION}, using MSELoss")
        loss_fn = nn.MSELoss()
    
    return loss_fn

# ============================================================================
# METRICS FUNCTIONS
# ============================================================================

def calculate_metrics(y_pred, y_true, metrics_list=None):
    """
    Calculate specified metrics.
    
    Args:
        y_pred: Predictions tensor
        y_true: True values tensor
        metrics_list: List of metric names (uses global METRICS if None)
    
    Returns:
        Dictionary of metric_name: value
    """
    if metrics_list is None:
        metrics_list = METRICS
    
    results = {}
    
    for metric in metrics_list:
        if metric == "MAE":
            results["MAE"] = torch.mean(torch.abs(y_pred - y_true)).item()
        elif metric == "MSE":
            results["MSE"] = torch.mean((y_pred - y_true) ** 2).item()
        elif metric == "R2":
            ss_res = torch.sum((y_true - y_pred) ** 2).item()
            ss_tot = torch.sum((y_true - torch.mean(y_true)) ** 2).item()
            r2 = 1 - (ss_res / (ss_tot + 1e-8))
            results["R2"] = r2
        elif metric == "RMSE":
            results["RMSE"] = torch.sqrt(torch.mean((y_pred - y_true) ** 2)).item()
    
    return results

# ============================================================================
# TRAINING FUNCTIONS
# ============================================================================

def train_epoch(model, train_loader, criterion, optimizer):
    """Train for one epoch."""
    model.train()
    total_loss = 0.0
    n_batches = 0
    
    for batch_X, batch_y in train_loader:
        batch_X = batch_X.to(DEVICE)
        batch_y = batch_y.to(DEVICE)
        
        # Forward pass
        optimizer.zero_grad()
        predictions = model(batch_X)
        loss = criterion(predictions, batch_y)
        
        # Backward pass
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
        n_batches += 1
    
    return total_loss / n_batches

def evaluate(model, data_loader, criterion):
    """Evaluate model on data."""
    model.eval()
    total_loss = 0.0
    all_predictions = []
    all_targets = []
    n_batches = 0
    
    with torch.no_grad():
        for batch_X, batch_y in data_loader:
            batch_X = batch_X.to(DEVICE)
            batch_y = batch_y.to(DEVICE)
            
            predictions = model(batch_X)
            loss = criterion(predictions, batch_y)
            
            total_loss += loss.item()
            all_predictions.append(predictions.cpu())
            all_targets.append(batch_y.cpu())
            n_batches += 1
    
    avg_loss = total_loss / n_batches
    all_predictions = torch.cat(all_predictions, dim=0)
    all_targets = torch.cat(all_targets, dim=0)
    
    metrics = calculate_metrics(all_predictions, all_targets)
    
    return avg_loss, metrics, all_predictions, all_targets

def train_model(model, train_loader, test_loader, criterion, optimizer):
    """Main training loop."""
    logger = logging.getLogger()
    logger.info("="*60)
    logger.info("Starting Training")
    logger.info("="*60)
    logger.info(f"Device: {DEVICE}")
    logger.info(f"Epochs: {EPOCHS}, Batch Size: {BATCH_SIZE}, Learning Rate: {LEARNING_RATE}")
    
    best_metric_value = float('inf') if METRICS_THRESHOLD_TYPE == "min" else float('-inf')
    best_model_state = None
    patience_counter = 0
    training_history = {
        'train_loss': [],
        'test_loss': [],
        'metrics': {metric: [] for metric in METRICS}
    }
    
    start_time = time.time()
    
    for epoch in range(1, EPOCHS + 1):
        # Train
        train_loss = train_epoch(model, train_loader, criterion, optimizer)
        
        # Evaluate
        test_loss, test_metrics, _, _ = evaluate(model, test_loader, criterion)
        
        # Store history
        training_history['train_loss'].append(train_loss)
        training_history['test_loss'].append(test_loss)
        for metric_name, metric_value in test_metrics.items():
            training_history['metrics'][metric_name].append(metric_value)
        
        # Logging
        if epoch % LOG_INTERVAL == 0 or epoch == 1:
            metric_str = ", ".join([f"{k}: {v:.6f}" for k, v in test_metrics.items()])
            logger.info(f"Epoch [{epoch:4d}/{EPOCHS}] | Train Loss: {train_loss:.6f} | "
                      f"Test Loss: {test_loss:.6f} | {metric_str}")
        
        # Save checkpoint
        if epoch % SAVE_INTERVAL == 0:
            save_model(model, epoch, test_loss, test_metrics)
        
        # Early stopping check
        if EARLY_STOPPING:
            metric_to_check = test_loss if EARLY_STOPPING_METRIC == "loss" else test_metrics.get(EARLY_STOPPING_METRIC, test_loss)
            
            if METRICS_THRESHOLD_TYPE == "min":
                improved = metric_to_check < (best_metric_value - EARLY_STOPPING_MIN_DELTA)
            else:
                improved = metric_to_check > (best_metric_value + EARLY_STOPPING_MIN_DELTA)
            
            if improved:
                best_metric_value = metric_to_check
                best_model_state = model.state_dict().copy()
                patience_counter = 0
            else:
                patience_counter += 1
                if patience_counter >= EARLY_STOPPING_PATIENCE:
                    logger.info(f"Early stopping triggered at epoch {epoch}")
                    if best_model_state is not None:
                        model.load_state_dict(best_model_state)
                    break
    
    elapsed_time = time.time() - start_time
    logger.info("="*60)
    logger.info(f"Training completed in {elapsed_time:.2f} seconds")
    logger.info("="*60)
    
    return training_history

# ============================================================================
# MODEL SAVING/LOADING
# ============================================================================

def save_model(model, epoch, loss, metrics, suffix=""):
    """
    Save model checkpoint.
    
    Note: torch.save() uses Python's pickle module by default. This is the standard
    PyTorch approach and works well for trusted files. However, pickle can execute
    arbitrary code when loading, so only load files from trusted sources.
    
    For untrusted files, consider:
    - Using SAVE_WEIGHTS_ONLY=True to save only state_dict
    - Using LOAD_WEIGHTS_ONLY=True when loading (requires PyTorch 1.13.0+)
    - Exporting to ONNX or TorchScript format for cross-platform deployment
    """
    if SAVE_WEIGHTS_ONLY:
        # Save only the state_dict (weights) - safer but requires model architecture
        checkpoint = model.state_dict()
        # Use model name if provided, otherwise use model version
        if MODEL_NAME:
            name_part = MODEL_NAME
        else:
            name_part = f"model_v{MODEL_VERSION}"
        filename = f"{name_part}_epoch{epoch}{suffix}_weights_only.{SAVE_FORMAT}"
    else:
        # Save full checkpoint with metadata
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'loss': loss,
            'metrics': metrics,
            'model_config': {
                'algorithm': ALGORITHM,
                'input_size': INPUT_SIZE,
                'hidden_sizes': HIDDEN_SIZES,
                'output_size': OUTPUT_SIZE,
                'activation': ACTIVATION
            },
            'training_config': {
                'optimizer': OPTIMIZER,
                'learning_rate': LEARNING_RATE,
                'loss_function': LOSS_FUNCTION
            }
        }
        # Use model name if provided, otherwise use model version
        if MODEL_NAME:
            name_part = MODEL_NAME
        else:
            name_part = f"model_v{MODEL_VERSION}"
        filename = f"{name_part}_epoch{epoch}{suffix}.{SAVE_FORMAT}"
    
    filepath = os.path.join(SAVE_PATH, filename)
    
    # Save with optional pickle protocol specification
    save_kwargs = {}
    if PICKLE_PROTOCOL is not None:
        save_kwargs['pickle_protocol'] = PICKLE_PROTOCOL
    
    torch.save(checkpoint, filepath, **save_kwargs)
    
    logger = logging.getLogger()
    logger.info(f"Model saved: {filepath} (format: {SAVE_FORMAT}, weights_only: {SAVE_WEIGHTS_ONLY})")

def load_model(model, filepath):
    """
    Load model checkpoint.
    
    WARNING: torch.load() uses pickle, which can execute arbitrary code.
    Only load files from trusted sources!
    
    For untrusted files:
    - Set LOAD_WEIGHTS_ONLY=True (requires PyTorch 1.13.0+)
    - This uses weights_only=True which only loads tensors, not arbitrary Python objects
    """
    load_kwargs = {'map_location': DEVICE}
    logger = logging.getLogger()
    
    # Use weights_only for safer loading (PyTorch 1.13.0+)
    if LOAD_WEIGHTS_ONLY:
        load_kwargs['weights_only'] = True
        logger.info("Loading with weights_only=True (safer, but may fail for older checkpoints)")
    
    # Try loading with weights_only, fall back if not supported
    try:
        checkpoint = torch.load(filepath, **load_kwargs)
    except TypeError as e:
        if LOAD_WEIGHTS_ONLY and 'weights_only' in str(e):
            logger.warning("weights_only parameter not supported in this PyTorch version. "
                          "Upgrade to PyTorch 1.13.0+ for safer loading. Loading without it...")
            load_kwargs.pop('weights_only', None)
            checkpoint = torch.load(filepath, **load_kwargs)
        else:
            raise
    
    # Handle both full checkpoint and weights-only formats
    if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
        # Full checkpoint format
        model.load_state_dict(checkpoint['model_state_dict'])
        logger.info(f"Model loaded from {filepath}")
        logger.info(f"Epoch: {checkpoint.get('epoch', 'N/A')}, Loss: {checkpoint.get('loss', 'N/A'):.6f}")
    else:
        # Weights-only format (just state_dict)
        model.load_state_dict(checkpoint)
        logger.info(f"Model weights loaded from {filepath} (weights-only format)")
    
    return checkpoint

# ============================================================================
# MAIN FUNCTION
# ============================================================================

def main():
    """Main training function."""
    # Setup
    setup_seed(SEED)
    setup_directories()
    logger = setup_logging()
    
    logger.info("="*60)
    logger.info("Training Template Started")
    logger.info("="*60)
    logger.info(f"Model Version: {MODEL_VERSION}")
    logger.info(f"Device: {DEVICE}")
    
    # ========================================================================
    # TODO: Load your data here
    # Replace this with your actual data loading
    # ========================================================================
    # Example: Generate dummy data (replace with your data loading)
    logger.info("Loading data...")
    # X_train, y_train = load_your_data()  # Replace this
    # For demonstration, creating dummy data:
    X_dummy = np.random.randn(1000, INPUT_SIZE).astype(np.float32)
    y_dummy = (2 * X_dummy.sum(axis=1) + 1 + np.random.randn(1000) * 0.1).astype(np.float32)
    y_dummy = y_dummy.reshape(-1, 1)
    
    # Prepare data
    train_loader, test_loader, data_stats = prepare_data(X_dummy, y_dummy)
    
    # Create model
    model = create_model()
    
    # Setup optimizer and loss
    optimizer = get_optimizer(model)
    criterion = get_loss_function()
    
    logger.info(f"Optimizer: {OPTIMIZER}, Loss: {LOSS_FUNCTION}")
    
    # Train
    history = train_model(model, train_loader, test_loader, criterion, optimizer)
    
    # Final evaluation
    logger.info("Performing final evaluation...")
    final_loss, final_metrics, predictions, targets = evaluate(model, test_loader, criterion)
    
    logger.info("="*60)
    logger.info("Final Results:")
    logger.info(f"Test Loss: {final_loss:.6f}")
    for metric_name, metric_value in final_metrics.items():
        logger.info(f"{metric_name}: {metric_value:.6f}")
    logger.info("="*60)
    
    # Save final model
    save_model(model, EPOCHS, final_loss, final_metrics, suffix="_final")
    
    logger.info("Training completed successfully!")

if __name__ == "__main__":
    main()