Upload AI Model Code

SAMFormer.ipynb

@@ -0,0 +1,904 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9c91ef58",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "import tensorflow as tf\n",
+    "import tensorflow.keras.backend as K\n",
+    "from tensorflow.keras import layers\n",
+    "\n",
+    "import collections\n",
+    "import logging\n",
+    "import numpy as np\n",
+    "import tensorflow as tf\n",
+    "from tensorflow.keras.callbacks import EarlyStopping, LearningRateScheduler, ModelCheckpoint, ReduceLROnPlateau\n",
+    "import pandas as pd\n",
+    "from sklearn.preprocessing import StandardScaler\n",
+    "import os\n",
+    "import math\n",
+    "import gc"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9a965750",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "name = \"ETTm1\" # dataset name\n",
+    "seq_len = 512\n",
+    "\n",
+    "batch_size = 32\n",
+    "pred_len = 192\n",
+    "\n",
+    "feature_type =  \"M\"\n",
+    "target = \"\"\n",
+    "learning_rate = 0.0001\n",
+    "patience = 5\n",
+    "\n",
+    "num_heads=1\n",
+    "d_model=16\n",
+    "rho = 0.1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d17cf421",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# please provide your own absolute links here\n",
+    "LOCAL_CACHE_DIR = '../Data/Benchmark/' # please include your dataset directory here\n",
+    "\n",
+    "checkpoint_path = \"../Train_SAM/models/checkpoint\" + name + str(pred_len) + \"_SAM\" +  \".model\"  # where the model checkpoint should be saved + the .model file type\n",
+    "name_df = \"../Results_SAM/\" + name + \"_\" +  str(pred_len) + \"_SAM\" + \".csv\"  # a filepath .csv file type to save the mse and mae results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5a6e0049",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class CaptureWeightsCallback(tf.keras.callbacks.Callback):\n",
+    "    \"\"\"\n",
+    "    Custom TensorFlow callback for capturing and logging model weights during training, with a focus on attention weights.\n",
+    "\n",
+    "    This callback is designed to monitor the evolution of model weights, particularly attention weights, across training epochs.\n",
+    "    It facilitates the analysis of training dynamics and model behavior by storing weight snapshots at specified intervals.\n",
+    "\n",
+    "    Attributes:\n",
+    "        model (tf.keras.Model): Instance of the TensorFlow model being trained. The model should have a method\n",
+    "                                `get_last_attention_weights()` that this callback can invoke to obtain attention weights.\n",
+    "        attention_weights_history (list): Accumulates the attention weights captured at the end of specified epochs. \n",
+    "                                          This history facilitates post-training analysis of weight adjustments.\n",
+    "\n",
+    "    Methods:\n",
+    "        on_epoch_end(epoch, logs=None): Overrides the base class method to capture attention weights at the end of each epoch.\n",
+    "                                        Weights are captured based on specified criteria, e.g., every 5 epochs.\n",
+    "        get_attention_weights_history(): Provides access to the accumulated history of attention weights captured during training.\n",
+    "    \"\"\"\n",
+    "    \n",
+    "    def __init__(self, model):\n",
+    "        \"\"\"\n",
+    "        Initializes the callback with a specific model to monitor its attention weights during training.\n",
+    "\n",
+    "        Parameters:\n",
+    "            model (tf.keras.Model): The model whose attention weights are to be monitored and captured.\n",
+    "        \"\"\"\n",
+    "        super().__init__()\n",
+    "        self.model = model\n",
+    "        self.penultimate_weights = None\n",
+    "        self.attention_weights_history = []\n",
+    "\n",
+    "    def on_epoch_end(self, epoch, logs=None):\n",
+    "        \"\"\"\n",
+    "        Called at the end of an epoch during training to capture and store attention weights if the current\n",
+    "        epoch satisfies the capture criteria (e.g., every 5 epochs).\n",
+    "\n",
+    "        Parameters:\n",
+    "            epoch (int): The current epoch number.\n",
+    "            logs (dict): Currently unused. Contains logs from the training epoch.\n",
+    "        \"\"\"\n",
+    "        if epoch % 5 == 0:  # Perform analysis every 5 epochs\n",
+    "            # Retrieve attention weights from the model\n",
+    "            last_attention_weights = self.model.get_last_attention_weights()\n",
+    "            if last_attention_weights is not None:\n",
+    "                self.attention_weights_history.append(last_attention_weights)\n",
+    "    \n",
+    "    def get_attention_weights_history(self):\n",
+    "        \"\"\"\n",
+    "        Returns the history of attention weights captured during training.\n",
+    "\n",
+    "        Returns:\n",
+    "            A list of attention weights captured at specified intervals during training.\n",
+    "        \"\"\"\n",
+    "        return self.attention_weights_history"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "22e2e8be",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def setup_callbacks(learning_rate, patience, checkpoint_path, model):\n",
+    "    \"\"\"\n",
+    "    Sets up and returns TensorFlow callbacks for use during model training. These callbacks include early stopping,\n",
+    "    learning rate scheduling, model checkpointing, and a custom callback for capturing model weights.\n",
+    "\n",
+    "    This function is tailored to support flexible training configurations, allowing for dynamic adjustment of training\n",
+    "    behavior based on model performance and training progress.\n",
+    "\n",
+    "    Parameters:\n",
+    "        args (argparse.Namespace): Parsed command-line arguments containing training configurations such as patience for early stopping,\n",
+    "                                   total training epochs, and initial learning rate.\n",
+    "        checkpoint_path (str): File path where model checkpoints will be saved. The best model according to validation loss is checkpointed.\n",
+    "        model (tf.keras.Model): The TensorFlow model being trained. Required for initializing the `CaptureWeightsCallback`.\n",
+    "        model_name (str): Name of the model being trained. This can be used to adjust callback behavior for different models.\n",
+    "\n",
+    "    Returns:\n",
+    "        tuple: A tuple containing:\n",
+    "               - A list of TensorFlow callbacks configured for the training session.\n",
+    "               - An instance of `CaptureWeightsCallback`, which can be used post-training to access captured weights.\n",
+    "\n",
+    "    Raises:\n",
+    "        Exception: If an error occurs in the setup of callbacks, an exception is logged and raised to prevent silent training failures.\n",
+    "\n",
+    "    Example:\n",
+    "        >>> callbacks, capture_weights_callback = setup_callbacks(args, './model_checkpoints', model, 'my_model')\n",
+    "        This example demonstrates how to invoke `setup_callbacks` to obtain configured callbacks for training, including a custom\n",
+    "        weight capture callback for post-training analysis.\n",
+    "    \"\"\"\n",
+    "    try:\n",
+    "        checkpoint_callback = ModelCheckpoint(\n",
+    "            filepath=checkpoint_path,\n",
+    "            monitor='val_loss',  \n",
+    "            verbose=1,\n",
+    "            save_best_only=True,\n",
+    "            save_weights_only=True,\n",
+    "        )\n",
+    "\n",
+    "        early_stop_callback = EarlyStopping(\n",
+    "            monitor='val_loss',\n",
+    "            patience=patience,\n",
+    "            verbose=1,\n",
+    "        )\n",
+    "\n",
+    "        lr_schedule_callback = LearningRateScheduler(\n",
+    "            lambda epoch: cosine_annealing(epoch, 5, learning_rate, 1e-6),\n",
+    "            verbose=1,\n",
+    "        )\n",
+    "\n",
+    "        lrdecay = ReduceLROnPlateau(monitor='val_loss', factor=0.90, patience=3, min_lr=learning_rate * 0.001, verbose=1)\n",
+    "\n",
+    "        capture_weights_callback = CaptureWeightsCallback(model)\n",
+    "\n",
+    "        callbacks = [checkpoint_callback, lr_schedule_callback, capture_weights_callback, early_stop_callback]\n",
+    "        \n",
+    "        return callbacks, capture_weights_callback\n",
+    "    except Exception as e:\n",
+    "        logging.error(f\"Error setting up callbacks: {e}\")\n",
+    "        raise\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5604ad15",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def cosine_annealing(epoch, max_epochs, initial_lr, min_lr):\n",
+    "    \"\"\"\n",
+    "    Applies cosine annealing to the learning rate.\n",
+    "    \n",
+    "    Parameters:\n",
+    "        epoch (int): Current epoch.\n",
+    "        max_epochs (int): Maximum number of epochs.\n",
+    "        initial_lr (float): Initial learning rate.\n",
+    "        min_lr (float): Minimum learning rate.\n",
+    "        \n",
+    "    Returns:\n",
+    "        float: Adjusted learning rate.\n",
+    "    \"\"\"\n",
+    "    cos_inner = (math.pi * (epoch % max_epochs)) / max_epochs\n",
+    "    return min_lr + (initial_lr - min_lr) * (math.cos(cos_inner) + 1) / 2\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "430e2b9e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class RevNorm(layers.Layer):\n",
+    "  \"\"\"\n",
+    "  Implements Reversible Instance Normalization (RevNorm).\n",
+    "\n",
+    "  This layer normalizes input features per instance and can reverse the normalization process. It is designed\n",
+    "  to maintain the statistical properties of the input data, making it particularly useful in generative models\n",
+    "  where the exact inverse operation is necessary.\n",
+    "\n",
+    "  Attributes:\n",
+    "      axis (int): The axis along which to compute the mean and standard deviation for normalization.\n",
+    "      eps (float): A small constant added to the standard deviation to prevent division by zero.\n",
+    "      affine (bool): Whether to apply a learnable affine transformation after normalization.\n",
+    "\n",
+    "  The original implementation can be found in the Google Research repository:\n",
+    "  https://github.com/google-research/google-research/blob/master/tsmixer/tsmixer_basic/models/rev_in.py\n",
+    "\n",
+    "  Example usage:\n",
+    "      rev_norm = RevNorm(axis=-1, eps=1e-5, affine=True)\n",
+    "      normalized_output = rev_norm(input_tensor, mode='norm')\n",
+    "      denormalized_output = rev_norm(normalized_output, mode='denorm', target_slice=slice_indices)\n",
+    "  \"\"\"\n",
+    "\n",
+    "  def __init__(self, axis, eps=1e-5, affine=True):\n",
+    "    super().__init__()\n",
+    "    self.axis = axis  # Defines the dimension along which normalization is performed.\n",
+    "    self.eps = eps  # Small epsilon value to ensure numerical stability.\n",
+    "    self.affine = affine  # Determines if learnable affine parameters should be used.\n",
+    "\n",
+    "  def build(self, input_shape):\n",
+    "    \"\"\"\n",
+    "    Initializes the layer's weights.\n",
+    "\n",
+    "    This method creates affine transformation weights if the `affine` attribute is set to True.\n",
+    "    It defines two trainable weights, `affine_weight` and `affine_bias`, which are used to scale and shift\n",
+    "    the normalized data respectively.\n",
+    "\n",
+    "    Args:\n",
+    "        input_shape (TensorShape): The shape of the input tensor to the layer. The last dimension is used\n",
+    "                                   to determine the shape of the affine weights.\n",
+    "\n",
+    "    Note: This method is automatically called during the first use of the layer.\n",
+    "    \"\"\"\n",
+    "    if self.affine:\n",
+    "      self.affine_weight = self.add_weight(\n",
+    "          'affine_weight', shape=input_shape[-1], initializer='ones'\n",
+    "      )\n",
+    "      self.affine_bias = self.add_weight(\n",
+    "          'affine_bias', shape=input_shape[-1], initializer='zeros'\n",
+    "      )\n",
+    "\n",
+    "  def call(self, x, mode, target_slice=None):\n",
+    "    \"\"\"\n",
+    "    Performs normalization or denormalization on the input tensor.\n",
+    "\n",
+    "    Args:\n",
+    "        x (Tensor): Input tensor to be normalized or denormalized.\n",
+    "        mode (str): 'norm' for normalization and 'denorm' for denormalization.\n",
+    "        target_slice (slice, optional): Target slice for denormalization.\n",
+    "\n",
+    "    Returns:\n",
+    "        Tensor: The normalized or denormalized output.\n",
+    "    \"\"\"\n",
+    "    if mode == 'norm':\n",
+    "      self._get_statistics(x)\n",
+    "      x = self._normalize(x)\n",
+    "    elif mode == 'denorm':\n",
+    "      x = self._denormalize(x, target_slice)\n",
+    "    else:\n",
+    "      raise NotImplementedError\n",
+    "    return x\n",
+    "\n",
+    "  def _get_statistics(self, x):\n",
+    "    \"\"\"\n",
+    "    Computes the mean and standard deviation of the input tensor along the specified axis.\n",
+    "\n",
+    "    The calculated mean and standard deviation are used for normalizing the input data. They are computed\n",
+    "    using `tf.reduce_mean` and `tf.sqrt(tf.reduce_variance(...) + self.eps)` to ensure numerical stability.\n",
+    "\n",
+    "    Args:\n",
+    "        x (Tensor): Input tensor from which the statistics are computed.\n",
+    "\n",
+    "    Updates:\n",
+    "        self.mean (Tensor): The mean of the input tensor, calculated along the specified axis.\n",
+    "        self.stdev (Tensor): The standard deviation of the input tensor, ensuring numerical stability by adding `self.eps`.\n",
+    "    \"\"\"\n",
+    "    self.mean = tf.stop_gradient(\n",
+    "        tf.reduce_mean(x, axis=self.axis, keepdims=True)\n",
+    "    )\n",
+    "    self.stdev = tf.stop_gradient(\n",
+    "        tf.sqrt(\n",
+    "            tf.math.reduce_variance(x, axis=self.axis, keepdims=True) + self.eps\n",
+    "        )\n",
+    "    )\n",
+    "\n",
+    "  def _normalize(self, x):\n",
+    "    \"\"\"\n",
+    "    Normalizes the input tensor using the computed mean and standard deviation.\n",
+    "\n",
+    "    This method subtracts the mean from the input tensor and divides it by the standard deviation, effectively\n",
+    "    standardizing the input to have a mean of 0 and a standard deviation of 1. If affine transformation is enabled,\n",
+    "    it further applies scaling and shifting to the standardized input.\n",
+    "\n",
+    "    Args:\n",
+    "        x (Tensor): Input tensor to be normalized.\n",
+    "\n",
+    "    Returns:\n",
+    "        Tensor: The normalized tensor.\n",
+    "    \"\"\"\n",
+    "    x = x - self.mean\n",
+    "    x = x / self.stdev\n",
+    "    if self.affine:\n",
+    "      x = x * self.affine_weight\n",
+    "      x = x + self.affine_bias\n",
+    "    return x\n",
+    "\n",
+    "  def _denormalize(self, x, target_slice=None):\n",
+    "    \"\"\"\n",
+    "    Reverses the normalization process for the given slice of the input tensor.\n",
+    "\n",
+    "    This method applies the inverse of the normalization operation. If affine transformation was applied during\n",
+    "    normalization, it reverses this process first. Then, it multiplies the tensor by the standard deviation and adds\n",
+    "    the mean to denormalize the data.\n",
+    "\n",
+    "    Args:\n",
+    "        x (Tensor): Normalized tensor that needs to be denormalized.\n",
+    "        target_slice (slice, optional): Specific slice of the tensor to denormalize. Useful when different parts\n",
+    "                                         of the tensor require different reverse operations.\n",
+    "\n",
+    "    Returns:\n",
+    "        Tensor: The denormalized tensor.\n",
+    "    \"\"\"\n",
+    "    if self.affine:\n",
+    "      x = x - self.affine_bias[target_slice]\n",
+    "      x = x / self.affine_weight[target_slice]\n",
+    "    x = x * self.stdev[:, :, target_slice]\n",
+    "    x = x + self.mean[:, :, target_slice]\n",
+    "    return x\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "323930c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class SAM:\n",
+    "    \"\"\"\n",
+    "    Sharpness-Aware Minimization (SAM) for Enhanced Training Stability.\n",
+    "    \n",
+    "    SAM optimizes a model's parameters in the direction that enhances model\n",
+    "    performance while simultaneously minimizing loss sharpness, aiming to improve\n",
+    "    generalization. This implementation wraps around a base TensorFlow optimizer\n",
+    "    to apply the SAM methodology.\n",
+    "    \n",
+    "    Reference:\n",
+    "    \"Sharpness-Aware Minimization for Efficiently Improving Generalization\"\n",
+    "    by Foret, Kleiner, Mobahi, and Neyshabur. https://openreview.net/pdf?id=6Tm1mposlrM\n",
+    "\n",
+    "    The original implementation can be found at:\n",
+    "    - https://github.com/davda54/sam\n",
+    "    \n",
+    "    Attributes:\n",
+    "        base_optimizer (tf.keras.optimizers.Optimizer): The TensorFlow optimizer to wrap.\n",
+    "        rho (float): The neighborhood size for sharpness-aware optimization.\n",
+    "        eps (float): A small epsilon value to prevent division by zero.\n",
+    "    \"\"\"\n",
+    "\n",
+    "    def __init__(self, base_optimizer, rho=0.05, eps=1e-12):\n",
+    "        \"\"\"\n",
+    "        Initializes the SAM optimizer wrapper.\n",
+    "        \n",
+    "        Parameters:\n",
+    "            base_optimizer (tf.keras.optimizers.Optimizer): The base optimizer.\n",
+    "            rho (float): The neighborhood size for sharpness-aware optimization.\n",
+    "            eps (float): A small epsilon value to prevent division by zero.\n",
+    "        \"\"\"\n",
+    "        assert rho >= 0.0, f\"Invalid rho, should be non-negative: {rho}\"\n",
+    "        self.rho = rho\n",
+    "        self.eps = eps\n",
+    "        self.base_optimizer = base_optimizer\n",
+    "\n",
+    "    def first_step(self, gradients, trainable_vars):\n",
+    "        \"\"\"\n",
+    "        Performs the first optimization step, moving weights in the direction\n",
+    "        that increases loss sharpness.\n",
+    "        \n",
+    "        Parameters:\n",
+    "            gradients (List[tf.Tensor]): Gradients of the loss with respect to the model parameters.\n",
+    "            trainable_vars (List[tf.Variable]): The model's trainable variables.\n",
+    "        \"\"\"\n",
+    "        self.e_ws = []\n",
+    "        grad_norm = tf.linalg.global_norm(gradients)\n",
+    "        ew_multiplier = self.rho / (grad_norm + self.eps)\n",
+    "\n",
+    "        for i in range(len(trainable_vars)):\n",
+    "            e_w = tf.math.multiply(gradients[i], ew_multiplier)\n",
+    "            trainable_vars[i].assign_add(e_w)\n",
+    "            self.e_ws.append(e_w)\n",
+    "\n",
+    "    def second_step(self, gradients, trainable_variables):\n",
+    "        \"\"\"\n",
+    "        Performs the second optimization step, applying the base optimizer\n",
+    "        update after reverting the first step's perturbation.\n",
+    "        \n",
+    "        Parameters:\n",
+    "            gradients (List[tf.Tensor]): Gradients of the loss with respect to the model parameters after the first step.\n",
+    "            trainable_variables (List[tf.Variable]): The model's trainable variables.\n",
+    "        \"\"\"\n",
+    "        for i in range(len(trainable_variables)):\n",
+    "            trainable_variables[i].assign_add(-self.e_ws[i])  # Revert first step\n",
+    "        self.base_optimizer.apply_gradients(zip(gradients, trainable_variables))\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3adf7700",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class SpectralNormalizedAttention(layers.MultiHeadAttention):\n",
+    "    \"\"\"\n",
+    "    Spectral Normalized Multi-Head Attention Layer.\n",
+    "    \n",
+    "    This layer extends the MultiHeadAttention layer with spectral normalization on the\n",
+    "    query, key, and value weights, implementing the sigma-reparam method described in\n",
+    "    \"Stabilizing Transformer Training by Preventing Attention Entropy Collapse\".\n",
+    "    \n",
+    "    Paper URL: https://openreview.net/forum?id=LL8gz8FHxH\n",
+    "    GitHub Code: https://github.com/apple/ml-sigma-reparam\n",
+    "    \n",
+    "    Attributes:\n",
+    "        gamma (tf.Variable): Scaling factor for the normalized weights, trainable.\n",
+    "    \"\"\"\n",
+    "\n",
+    "    def __init__(self, *args, **kwargs):\n",
+    "        \"\"\"Initializes the SpectralNormalizedAttention layer with standard arguments for MultiHeadAttention.\"\"\"\n",
+    "        super(SpectralNormalizedAttention, self).__init__(*args, **kwargs)\n",
+    "        self.gamma = self.add_weight(name='gamma', shape=[], initializer='ones', trainable=True)\n",
+    "\n",
+    "    def build(self, input_shape):\n",
+    "        \"\"\"Builds the layer, initializing weights.\"\"\"\n",
+    "        super(SpectralNormalizedAttention, self).build(input_shape)\n",
+    "        # Additional initializations can be added here if necessary.\n",
+    "\n",
+    "    def _normalize_weights(self, W):\n",
+    "        \"\"\"\n",
+    "        Normalizes the weights matrix W using its spectral norm.\n",
+    "        \n",
+    "        Parameters:\n",
+    "            W (tf.Tensor): The weight matrix to normalize.\n",
+    "        \n",
+    "        Returns:\n",
+    "            tf.Tensor: Spectrally normalized weights.\n",
+    "        \"\"\"\n",
+    "        singular_values = tf.linalg.svd(W, compute_uv=False)\n",
+    "        spectral_norm = tf.reduce_max(singular_values)\n",
+    "        return W / spectral_norm\n",
+    "\n",
+    "    def call(self, query, value, key=None, attention_mask=None, return_attention_scores=False):\n",
+    "        \"\"\"\n",
+    "        Calls the SpectralNormalizedAttention layer. Normalizes the query, key, and value weights\n",
+    "        before calling the parent MultiHeadAttention layer.\n",
+    "        \n",
+    "        Parameters:\n",
+    "            query (tf.Tensor): Query tensor.\n",
+    "            value (tf.Tensor): Value tensor.\n",
+    "            key (tf.Tensor): Key tensor. Defaults to None, in which case the query is used as the key.\n",
+    "            attention_mask (tf.Tensor): Optional tensor to mask out certain positions from attending to others.\n",
+    "            return_attention_scores (bool): Flag to return attention scores along with output.\n",
+    "        \n",
+    "        Returns:\n",
+    "            A tuple of (output tensor, attention scores) if return_attention_scores is True, otherwise just the output tensor.\n",
+    "        \"\"\"\n",
+    "        key = query if key is None else key\n",
+    "        query = self._normalize_weights(query) * self.gamma\n",
+    "        key = self._normalize_weights(key) * self.gamma\n",
+    "        value = self._normalize_weights(value) * self.gamma\n",
+    "\n",
+    "        return super(SpectralNormalizedAttention, self).call(query, value, key, attention_mask, return_attention_scores)\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "663d4c7e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class BaseModel(tf.keras.Model):\n",
+    "     \"\"\"\n",
+    "    A base model class that integrates various enhancements including \n",
+    "    Reversible Instance Normalization and Channel-Wise Attention, and optionally,\n",
+    "    spectral normalization and SAM optimization. To use SAMformer, enable use_sam,\n",
+    "    use_attention, use_revin and trainable.\n",
+    "    \n",
+    "    Attributes:\n",
+    "        pred_len (int): The length of the output predictions.\n",
+    "        num_heads (int): The number of heads in the multi-head attention mechanism. \n",
+    "        d_model (int): The dimensionality of the embedding vectors.\n",
+    "        use_sam (bool): If True, applies Sharpness-Aware Minimization (SAM) optimization technique during training, \n",
+    "                        aiming to improve model generalization by considering the loss landscape's sharpness.\n",
+    "        use_attention (bool): If True, enables the multi-head attention mechanism in the model. If False, the model is\n",
+    "                              equivalent to a simple linear layer.\n",
+    "        use_revin (bool): If True, applies Reversible Instance Normalization (RevIN) to the model.\n",
+    "        trainable (bool): Specifies if the model's weights should be updated or frozen during training. Useful to \n",
+    "                          highlight some attention layer issues in Time Series Forecasting.\n",
+    "        rho (float): The neighborhood size parameter for SAM optimization. It determines the radius within which SAM \n",
+    "                     seeks to minimize the sharpness of the loss landscape.\n",
+    "        spec (bool): If True, applies spectral normalization (sigma-reparam) to the attention mechanism, aiming to\n",
+    "                     stabilize the training by constraining the spectral norm of the weight matrices.\n",
+    "        \n",
+    "    Methods:\n",
+    "        call(inputs, training=False): Defines the computation from inputs to outputs, optionally applying SAM, \n",
+    "                                      spectral normalization, and reversible instance normalization based on the \n",
+    "                                      configuration.\n",
+    "        _apply_attention(x): Applies the attention mechanism to the input tensor, capturing the inter-dependencies \n",
+    "                             within the data thanks to the Channel-Wise Attention mechanism.\n",
+    "        get_last_attention_weights(): Retrieves the attention weights from the last but one batch, useful for \n",
+    "                                      analysis and debugging purposes.\n",
+    "        train_step(data): Custom training logic, including the application of SAM's two-step optimization process, \n",
+    "                          to improve model generalization and performance stability.\n",
+    "\n",
+    "    \"\"\"\n",
+    "\n",
+    "     def __init__(self, pred_len, num_heads=1, d_model=16, use_sam=None, \n",
+    "                 use_attention=None, use_revin=None, \n",
+    "                 trainable=None, rho=None, spec=None):\n",
+    "        super(BaseModel, self).__init__()\n",
+    "        self.pred_len = pred_len\n",
+    "        self.num_heads = num_heads\n",
+    "        self.d_model = d_model\n",
+    "        self.use_sam = use_sam\n",
+    "        self.use_attention = use_attention\n",
+    "        self.use_revin = use_revin\n",
+    "        self.rho = rho if use_sam and trainable else 0.0\n",
+    "        self.spec = spec\n",
+    "\n",
+    "        # Define model layers\n",
+    "        self.rev_norm = RevNorm(axis=-2)\n",
+    "        self.attention_layer = layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model)\n",
+    "        self.dense = layers.Dense(pred_len)\n",
+    "        self.all_attention_weights = collections.deque(maxlen=2)\n",
+    "        self.all_dense_weights = collections.deque(maxlen=2)\n",
+    "\n",
+    "        if self.spec:\n",
+    "            self.spec_layer = SpectralNormalizedAttention(num_heads=num_heads, key_dim=d_model)\n",
+    "\n",
+    "        #Define trainability of attention layer\n",
+    "        self.attention_layer.trainable = trainable\n",
+    "\n",
+    "     def call(self, inputs, training=False):\n",
+    "        \"\"\"\n",
+    "        The forward pass for the model.\n",
+    "        \n",
+    "        Parameters:\n",
+    "            inputs (Tensor): Input tensor.\n",
+    "            training (bool): Whether the call is for training.\n",
+    "        \n",
+    "        Returns:\n",
+    "            Tensor: The output of the model.\n",
+    "        \"\"\"\n",
+    "\n",
+    "        x = inputs\n",
+    "        if self.use_revin:\n",
+    "            x = self.rev_norm(x, mode='norm')\n",
+    "        x = tf.transpose(x, perm=[0, 2, 1])\n",
+    "\n",
+    "        if self.use_attention:\n",
+    "            attention_output = self._apply_attention(x)\n",
+    "            x = layers.Add()([x, attention_output])\n",
+    "\n",
+    "        x = self.dense(x)\n",
+    "        outputs = tf.transpose(x, perm=[0, 2, 1])\n",
+    "\n",
+    "        if self.use_revin:\n",
+    "            outputs = self.rev_norm(outputs, mode='denorm')\n",
+    "\n",
+    "        return outputs\n",
+    "\n",
+    "     def _apply_attention(self, x):\n",
+    "        \"\"\"\n",
+    "        Applies the attention mechanism to the input tensor.\n",
+    "        \n",
+    "        Parameters:\n",
+    "            x (Tensor): The input tensor.\n",
+    "            training (bool): Whether the call is for training.\n",
+    "        \n",
+    "        Returns:\n",
+    "            Tensor: The output tensor after applying attention.\n",
+    "        \"\"\"\n",
+    "        if self.spec:\n",
+    "            attention_output, weights = self.spec_layer(x, x, return_attention_scores=True)\n",
+    "        else:\n",
+    "            attention_output, weights = self.attention_layer(x, x, return_attention_scores=True)\n",
+    "        \n",
+    "        self.all_attention_weights.append(weights.numpy())\n",
+    "        return attention_output\n",
+    "\n",
+    "     def get_last_attention_weights(self):\n",
+    "        \"\"\"Returns the attention weights from the last but one batch.\"\"\"\n",
+    "        if len(self.all_attention_weights) > 1:\n",
+    "            return self.all_attention_weights[-2]\n",
+    "        return None\n",
+    "\n",
+    "     def train_step(self, data):\n",
+    "        sam_optimizer = SAM(self.optimizer, rho=self.rho, eps=1e-12) \n",
+    "\n",
+    "        # Unpack the data.\n",
+    "        x, y = data\n",
+    "\n",
+    "        with tf.GradientTape() as tape:\n",
+    "            y_pred = self(x, training=True)  # Forward pass\n",
+    "            loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)\n",
+    "\n",
+    "        # Compute gradients\n",
+    "        gradients = tape.gradient(loss, self.trainable_variables)\n",
+    "\n",
+    "        # Apply SAM's first step\n",
+    "        sam_optimizer.first_step(gradients, self.trainable_variables)\n",
+    "\n",
+    "        with tf.GradientTape() as tape:\n",
+    "            y_pred = self(x, training=True)  # Forward pass again\n",
+    "            loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)\n",
+    "\n",
+    "        # Compute gradients again\n",
+    "        gradients = tape.gradient(loss, self.trainable_variables)\n",
+    "\n",
+    "        # Apply SAM's second step\n",
+    "        sam_optimizer.second_step(gradients, self.trainable_variables)\n",
+    "\n",
+    "        # Update metrics\n",
+    "        self.compiled_metrics.update_state(y, y_pred)\n",
+    "\n",
+    "        # Return a dict mapping metric names to current value\n",
+    "        return {m.name: m.result() for m in self.metrics}\n",
+    "        "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cdc39fe9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class TSFDataLoader:\n",
+    "  \"\"\"Generate data loader from raw data.\"\"\"\n",
+    "\n",
+    "  def __init__(\n",
+    "      self, data, batch_size, seq_len, pred_len, feature_type, target='OT'\n",
+    "  ):\n",
+    "    self.data = data\n",
+    "    self.batch_size = batch_size\n",
+    "    self.seq_len = seq_len\n",
+    "    self.pred_len = pred_len\n",
+    "    self.feature_type = feature_type\n",
+    "    self.target = target\n",
+    "    self.target_slice = slice(0, None)\n",
+    "\n",
+    "    self._read_data()\n",
+    "\n",
+    "  def _read_data(self):\n",
+    "    \"\"\"Load raw data and split datasets.\"\"\"\n",
+    "\n",
+    "    # copy data from cloud storage if not exists\n",
+    "    if not os.path.isdir(LOCAL_CACHE_DIR):\n",
+    "      os.mkdir(LOCAL_CACHE_DIR)\n",
+    "\n",
+    "    file_name = self.data + '.csv'\n",
+    "    cache_filepath = os.path.join(LOCAL_CACHE_DIR, file_name)\n",
+    "    if not os.path.isfile(cache_filepath):\n",
+    "      tf.io.gfile.copy(\n",
+    "          os.path.join(DATA_DIR, file_name), cache_filepath, overwrite=True\n",
+    "      )\n",
+    "    df_raw = pd.read_csv(cache_filepath)\n",
+    "\n",
+    "    \n",
+    "    \n",
+    "    # S: univariate-univariate, M: multivariate-multivariate, MS:\n",
+    "    # multivariate-univariate\n",
+    "    df = df_raw.set_index('date')\n",
+    "    if self.feature_type == 'S':\n",
+    "      df = df[[self.target]]\n",
+    "    elif self.feature_type == 'MS':\n",
+    "      target_idx = df.columns.get_loc(self.target)\n",
+    "      self.target_slice = slice(target_idx, target_idx + 1)\n",
+    "\n",
+    "    # split train/valid/test\n",
+    "    n = len(df)\n",
+    "    if self.data.startswith('ETTm'):\n",
+    "      train_end = 12 * 30 * 24 * 4\n",
+    "      val_end = train_end + 4 * 30 * 24 * 4\n",
+    "      test_end = val_end + 4 * 30 * 24 * 4\n",
+    "    elif self.data.startswith('ETTh'):\n",
+    "      train_end = 12 * 30 * 24\n",
+    "      val_end = train_end + 4 * 30 * 24\n",
+    "      test_end = val_end + 4 * 30 * 24\n",
+    "    else:\n",
+    "      train_end = int(n * 0.7)\n",
+    "      val_end = n - int(n * 0.2)\n",
+    "      test_end = n\n",
+    "        \n",
+    "    train_df = df[:train_end]\n",
+    "    val_df = df[train_end - self.seq_len : val_end]\n",
+    "    test_df = df[val_end - self.seq_len : test_end]\n",
+    "\n",
+    "    # standardize by training set\n",
+    "    self.scaler = StandardScaler()\n",
+    "    self.scaler1 = StandardScaler()\n",
+    "      \n",
+    "    self.scaler.fit(train_df.values)\n",
+    "    # self.scaler1.fit(df.iloc[:train_end, :-5].values)\n",
+    "\n",
+    "    def scale_df(df, scaler):\n",
+    "      data = scaler.transform(df.values)\n",
+    "      # return pd.DataFrame(df, index=df.index, columns=df.columns)\n",
+    "      return pd.DataFrame(data, index=df.index, columns=df.columns)\n",
+    "\n",
+    "    self.train_df = scale_df(train_df, self.scaler)\n",
+    "    self.val_df = scale_df(val_df, self.scaler)\n",
+    "    self.test_df = scale_df(test_df, self.scaler)\n",
+    "    self.n_feature = self.train_df.shape[-1]\n",
+    "\n",
+    "  def _split_window(self, data):\n",
+    "    inputs = data[:, : self.seq_len, :]\n",
+    "    labels = data[:, self.seq_len :, self.target_slice]\n",
+    "    # Slicing doesn't preserve static shape information, so set the shapes\n",
+    "    # manually. This way the `tf.data.Datasets` are easier to inspect.\n",
+    "    inputs.set_shape([None, self.seq_len, None])\n",
+    "    labels.set_shape([None, self.pred_len, self.n_feature])\n",
+    "    return inputs, labels\n",
+    "\n",
+    "  def _make_dataset(self, data, shuffle=True):\n",
+    "    data = np.array(data, dtype=np.float32)\n",
+    "    ds = tf.keras.utils.timeseries_dataset_from_array(\n",
+    "        data=data,\n",
+    "        targets=None,\n",
+    "        sequence_length=(self.seq_len + self.pred_len),\n",
+    "        sequence_stride=1,\n",
+    "        shuffle=False,\n",
+    "        batch_size=self.batch_size,\n",
+    "    )\n",
+    "    ds = ds.map(self._split_window)\n",
+    "    return ds\n",
+    "\n",
+    "  def _make_dataset_test(self, data, shuffle=True):\n",
+    "    data = np.array(data, dtype=np.float32)\n",
+    "    ds = tf.keras.utils.timeseries_dataset_from_array(\n",
+    "             data=data,\n",
+    "             targets=None,\n",
+    "             sequence_length=(self.seq_len + self.pred_len),\n",
+    "             sequence_stride=1,\n",
+    "             shuffle=False,\n",
+    "             batch_size=1)\n",
+    "    ds = ds.map(self._split_window)\n",
+    "    return ds\n",
+    "\n",
+    "\n",
+    "  def inverse_transform(self, data):\n",
+    "    return self.scaler.inverse_transform(data)\n",
+    "\n",
+    "  def get_train(self, shuffle=True):\n",
+    "    return self._make_dataset(self.train_df, shuffle=shuffle)\n",
+    "\n",
+    "  def get_val(self):\n",
+    "    return self._make_dataset(self.val_df, shuffle=False)\n",
+    "\n",
+    "  def get_test(self):\n",
+    "    return self._make_dataset(self.test_df, shuffle=False)\n",
+    "\n",
+    "\n",
+    "def load_data(name, batch_size, seq_len, pred_len, feature_type, target):\n",
+    "    \"\"\"\n",
+    "    Loads or generates training, validation, and testing datasets based on the specified configurations.\n",
+    "\n",
+    "    Parameters:\n",
+    "        args (argparse.Namespace): Command line arguments specifying dataset configurations.\n",
+    "\n",
+    "    Returns:\n",
+    "        tuple: Training, validation, and test datasets as tf.data.Dataset objects.\n",
+    "    \"\"\"\n",
+    "    data_loader = TSFDataLoader(name, batch_size, seq_len, pred_len, feature_type, target)\n",
+    "    train_data, val_data, test_data = data_loader.get_train(), data_loader.get_val(), data_loader.get_test()\n",
+    "\n",
+    "    return train_data, val_data, test_data, data_loader.n_feature"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fe7abdc1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_data, val_data, test_data, n_features  = load_data(name, batch_size, seq_len, pred_len, feature_type, target)\n",
+    "new_cols = ['mse', 'mae', \"seq_len\", \"learning_rate\", \"batch_size\", \"num_heads\", \n",
+    "            \"d_model\", \"minimization_neighbourhood\"]\n",
+    "\n",
+    "df = pd.DataFrame(columns=new_cols)\n",
+    "\n",
+    "numEvalSamples = 1\n",
+    "seeds = [2022]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "29470575",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for sd in range(numEvalSamples):\n",
+    "    tf.keras.backend.clear_session() \n",
+    "\n",
+    "    tf.keras.utils.set_random_seed(seeds[sd])\n",
+    "    \n",
+    "    print(f\"..............iteration{sd}.................\")\n",
+    "    print()\n",
+    "    model = BaseModel(pred_len, num_heads=num_heads, d_model=d_model, use_sam=True, use_attention=True, use_revin=True, trainable=True, rho=rho, \n",
+    "                      spec=False)\n",
+    "    optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)\n",
+    "    model.compile(optimizer=optimizer, loss='mae', metrics=['mse', 'mae'], run_eagerly=True)\n",
+    "    \n",
+    "    callbacks, capture_weights_callback = setup_callbacks(learning_rate, patience,  checkpoint_path, model)\n",
+    "    \n",
+    "    history = model.fit(\n",
+    "                train_data,\n",
+    "                epochs=300,\n",
+    "                validation_data=val_data,\n",
+    "                callbacks=callbacks)\n",
+    "    \n",
+    "    model.load_weights(checkpoint_path)\n",
+    "    df.loc[sd, :] = model.evaluate(test_data)[-2:] + [ seq_len, learning_rate, batch_size, num_heads, \n",
+    "                                                d_model, rho]\n",
+    "    df.to_csv(name_df, header=True, index=False)\n",
+    "\n",
+    "    del model\n",
+    "    tf.keras.backend.clear_session()\n",
+    "    # garbage collector\n",
+    "    gc.collect()\n",
+    "\n",
+    "    # remove_dir(checkpoint_path)\n",
+    "\n",
+    "print(\"................completed Successfully......................\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "852cd02e",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}