Create Temporal_Convolutional_Network.py

Temporal_Convolutional_Network.py

@@ -0,0 +1,292 @@
+import os
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from sklearn.preprocessing import MinMaxScaler
+from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
+import matplotlib.pyplot as plt
+from itertools import product
+import json
+
+# -------------------------
+# Dataset
+# -------------------------
+class StockDataset(Dataset):
+    """Custom Dataset for stock price time-series forecasting."""
+    def __init__(self, series, seq_length):
+        self.series = series
+        self.seq_length = seq_length
+
+    def __len__(self):
+        return len(self.series) - self.seq_length
+
+    def __getitem__(self, idx):
+        x = self.series[idx:idx + self.seq_length]  # Shape: (seq_length,)
+        y = self.series[idx + self.seq_length]      # Shape: scalar
+        x = np.expand_dims(x, axis=0)               # Shape: (1, seq_length)
+        return torch.tensor(x, dtype=torch.float32), torch.tensor(y, dtype=torch.float32)
+
+# -------------------------
+# TCN Blocks
+# -------------------------
+class TemporalBlock(nn.Module):
+    """Temporal Convolutional Network block with causal dilated convolutions."""
+    def __init__(self, in_channels, out_channels, kernel_size, stride, dilation, dropout=0.2):
+        super().__init__()
+        padding = (kernel_size - 1) * dilation
+        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size,
+                               stride=stride, padding=padding, dilation=dilation)
+        self.relu1 = nn.ReLU()
+        self.dropout1 = nn.Dropout(dropout)
+        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size,
+                               stride=stride, padding=padding, dilation=dilation)
+        self.relu2 = nn.ReLU()
+        self.dropout2 = nn.Dropout(dropout)
+        self.downsample = nn.Conv1d(in_channels, out_channels, 1) if in_channels != out_channels else None
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = out[:, :, :x.size(2)]  # Trim padding
+        out = self.relu1(out)
+        out = self.dropout1(out)
+        out = self.conv2(out)
+        out = out[:, :, :x.size(2)]  # Trim padding
+        out = self.relu2(out)
+        out = self.dropout2(out)
+        res = x if self.downsample is None else self.downsample(x)
+        return self.relu(out + res)
+
+class TCN(nn.Module):
+    """Temporal Convolutional Network for time-series forecasting."""
+    def __init__(self, input_size, output_size, num_channels, kernel_size=3, dropout=0.2):
+        super().__init__()
+        layers = []
+        num_levels = len(num_channels)
+        for i in range(num_levels):
+            dilation_size = 2 ** i
+            in_channels = input_size if i == 0 else num_channels[i - 1]
+            out_channels = num_channels[i]
+            layers.append(
+                TemporalBlock(in_channels, out_channels, kernel_size,
+                              stride=1, dilation=dilation_size, dropout=dropout)
+            )
+        self.network = nn.Sequential(*layers)
+        self.linear = nn.Linear(num_channels[-1], output_size)
+
+    def forward(self, x):
+        out = self.network(x)
+        out = out[:, :, -1]
+        return self.linear(out)
+
+# -------------------------
+# Forecaster
+# -------------------------
+class StockPriceForecaster:
+    """Stock price forecasting with TCN model."""
+    def __init__(self, dataset_path, seq_length=30, batch_size=32, lr=0.001, epochs=20,
+                 kernel_size=3, num_channels=[32, 64, 64], dropout=0.2, test_split=0.2):
+        self.dataset_path = dataset_path
+        self.seq_length = seq_length
+        self.batch_size = batch_size
+        self.lr = lr
+        self.epochs = epochs
+        self.kernel_size = kernel_size
+        self.num_channels = num_channels
+        self.dropout = dropout
+        self.test_split = test_split
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.scaler = MinMaxScaler()
+
+    def load_data(self):
+        """Load and preprocess stock price data."""
+        if not os.path.exists(self.dataset_path):
+            raise FileNotFoundError(f"Dataset file not found at: {self.dataset_path}")
+        df = pd.read_csv(self.dataset_path)
+        if "Close" not in df.columns:
+            raise ValueError("CSV file must contain a 'Close' column")
+        prices = df["Close"].values.reshape(-1, 1)
+        prices_scaled = self.scaler.fit_transform(prices).flatten()
+        dataset = StockDataset(prices_scaled, self.seq_length)
+        train_size = int(len(dataset) * (1 - self.test_split))
+        test_size = len(dataset) - train_size
+        train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])
+        train_loader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True)
+        test_loader = DataLoader(test_dataset, batch_size=self.batch_size, shuffle=False)
+        return train_loader, test_loader
+
+    def train(self, model, train_loader):
+        """Train the TCN model."""
+        criterion = nn.MSELoss()
+        optimizer = torch.optim.Adam(model.parameters(), lr=self.lr)
+        model.train()
+        for epoch in range(self.epochs):
+            epoch_loss = 0
+            for x, y in train_loader:
+                x, y = x.to(self.device), y.to(self.device)
+                optimizer.zero_grad()
+                output = model(x)
+                loss = criterion(output.squeeze(), y)
+                loss.backward()
+                optimizer.step()
+                epoch_loss += loss.item()
+            print(f"Epoch [{epoch+1}/{self.epochs}], Loss: {epoch_loss/len(train_loader):.6f}")
+        return model
+
+    def evaluate(self, model, test_loader):
+        """Evaluate the model on the test set."""
+        model.eval()
+        predictions, actuals = [], []
+        with torch.no_grad():
+            for x, y in test_loader:
+                x, y = x.to(self.device), y.to(self.device)
+                output = model(x)
+                predictions.extend(output.squeeze().cpu().numpy())
+                actuals.extend(y.cpu().numpy())
+        predictions = self.scaler.inverse_transform(np.array(predictions).reshape(-1, 1)).flatten()
+        actuals = self.scaler.inverse_transform(np.array(actuals).reshape(-1, 1)).flatten()
+        mae = mean_absolute_error(actuals, predictions)
+        rmse = mean_squared_error(actuals, predictions, squared=False)
+        mape = np.mean(np.abs((actuals - predictions) / (actuals + 1e-10))) * 100
+        r2 = r2_score(actuals, predictions)
+        return mae, rmse, mape, r2, actuals, predictions
+
+    def run(self):
+        """Run training and evaluation."""
+        train_loader, test_loader = self.load_data()
+        model = TCN(input_size=1, output_size=1,
+                    num_channels=self.num_channels,
+                    kernel_size=self.kernel_size,
+                    dropout=self.dropout).to(self.device)
+        trained_model = self.train(model, train_loader)
+        return trained_model, self.evaluate(model, test_loader)
+
+# -------------------------
+# Save Model for Hugging Face
+# -------------------------
+def save_model_for_huggingface(model, scaler, config, save_dir="tcn_stock_model"):
+    """Save the model and necessary components for Hugging Face deployment."""
+    os.makedirs(save_dir, exist_ok=True)
+    
+    # Save model weights
+    torch.save(model.state_dict(), os.path.join(save_dir, "pytorch_model.bin"))
+    
+    # Save model configuration
+    with open(os.path.join(save_dir, "config.json"), "w") as f:
+        json.dump({
+            "input_size": 1,
+            "output_size": 1,
+            "num_channels": config["num_channels"],
+            "kernel_size": config["kernel_size"],
+            "dropout": config["dropout"],
+            "seq_length": config["seq_length"]
+        }, f, indent=4)
+    
+    # Save scaler for preprocessing
+    import pickle
+    with open(os.path.join(save_dir, "scaler.pkl"), "wb") as f:
+        pickle.dump(scaler, f)
+    
+    print(f"Model saved to {save_dir}")
+
+# -------------------------
+# Experiment Loop
+# -------------------------
+if __name__ == "__main__":
+    dataset_path = "/work/GOOGL.csv"  # Update to your CSV path
+
+    # Hyperparameter grid
+    seq_lengths = [20, 50]
+    batch_sizes = [16, 32]
+    learning_rates = [0.001, 0.0005]
+    kernel_sizes = [3, 5]
+    num_channels_list = [[32, 64, 128], [64, 128, 256]]
+    dropouts = [0.1, 0.2]
+
+    results = []
+    best_result = None
+    best_metrics = float('inf')  # Track best RMSE
+    best_model = None
+    best_config = None
+
+    # Run experiments
+    for seq, batch, lr, kernel, channels, dropout in product(
+        seq_lengths, batch_sizes, learning_rates, kernel_sizes, num_channels_list, dropouts
+    ):
+        print(f"\nRunning: seq={seq}, batch={batch}, lr={lr}, kernel={kernel}, channels={channels}, dropout={dropout}")
+        try:
+            forecaster = StockPriceForecaster(
+                dataset_path=dataset_path,
+                seq_length=seq,
+                batch_size=batch,
+                lr=lr,
+                epochs=20,
+                kernel_size=kernel,
+                num_channels=channels,
+                dropout=dropout,
+                test_split=0.2
+            )
+            model, (mae, rmse, mape, r2, actuals, predictions) = forecaster.run()
+            results.append({
+                "seq_length": seq,
+                "batch_size": batch,
+                "lr": lr,
+                "kernel_size": kernel,
+                "num_channels": str(channels),
+                "dropout": dropout,
+                "MAE": mae,
+                "RMSE": rmse,
+                "MAPE": mape,
+                "R2": r2
+            })
+            if rmse < best_metrics:
+                best_metrics = rmse
+                best_result = (actuals, predictions, seq, batch, lr, kernel, channels, dropout)
+                best_model = model
+                best_config = {
+                    "seq_length": seq,
+                    "batch_size": batch,
+                    "lr": lr,
+                    "kernel_size": kernel,
+                    "num_channels": channels,
+                    "dropout": dropout
+                }
+        except Exception as e:
+            print(f"Error with config seq={seq}, batch={batch}, lr={lr}, kernel={kernel}, channels={channels}, dropout={dropout}: {e}")
+            continue
+
+    # Save results
+    df_results = pd.DataFrame(results)
+    df_results.to_csv("tcn_experiments_results.csv", index=False)
+    print("\nAll experiments done! Results saved to 'tcn_experiments_results.csv'")
+
+    # Display metrics table
+    print("\nMetrics Table:")
+    pd.set_option('display.max_columns', None)
+    pd.set_option('display.width', 1000)
+    pd.set_option('display.float_format', '{:.6f}'.format)
+    print(df_results)
+
+    # Save best model for Hugging Face
+    if best_model is not None:
+        save_model_for_huggingface(best_model, forecaster.scaler, best_config)
+        print(f"\nBest model saved with RMSE: {best_metrics:.6f}")
+        print("\nBest configuration:")
+        print(pd.Series(best_config))
+
+    # Plot best combination
+    if best_result is not None:
+        actuals, predictions, seq, batch, lr, kernel, channels, dropout = best_result
+        plt.figure(figsize=(12, 6))
+        plt.plot(actuals, label="Actual Prices")
+        plt.plot(predictions, label="Predicted Prices")
+        plt.title(f"Best Model: seq={seq}, batch={batch}, lr={lr}, kernel={kernel}, channels={channels}, dropout={dropout}")
+        plt.xlabel("Time Step")
+        plt.ylabel("Price")
+        plt.legend()
+        plt.grid(True)
+        plt.show()
+    else:
+        print("No successful experiments to plot.")