{"metadata":{"kernelspec":{"display_name":"Python 3","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.10.14"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":84493,"databundleVersionId":9871156,"sourceType":"competition"},{"sourceId":203900450,"sourceType":"kernelVersion"}],"dockerImageVersionId":30787,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":7.594014,"end_time":"2024-10-10T11:58:36.355301","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2024-10-10T11:58:28.761287","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Useful notebooks:\n\n- Preprocessing : https://www.kaggle.com/code/motono0223/js24-preprocessing-create-lags\n- Training (XGB) : https://www.kaggle.com/code/motono0223/js24-train-gbdt-model-with-lags-singlemodel\n  - trained XGB model : https://www.kaggle.com/datasets/motono0223/js24-trained-gbdt-model\n- Training (NN): **this notebook** https://www.kaggle.com/code/voix97/jane-street-rmf-training-nn\n  - trained NN model : https://www.kaggle.com/datasets/voix97/js-xs-nn-trained-model\n- Inference of NN : https://www.kaggle.com/code/voix97/jane-street-rmf-nn-with-pytorch-lightning\n- Inference of NN+XGB:  https://www.kaggle.com/code/voix97/jane-street-rmf-nn-xgb\n- EDA(1) : https://www.kaggle.com/code/motono0223/eda-jane-street-real-time-market-data-forecasting\n- EDA(2) : https://www.kaggle.com/code/motono0223/eda-v2-jane-street-real-time-market-forecasting","metadata":{}},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Training Neural Networks (MLP) with PyTorch Lightning","metadata":{}},{"cell_type":"code","source":"import os\nimport pickle\nimport polars as pl\nimport numpy as np\nimport pandas as pd","metadata":{"execution":{"iopub.status.busy":"2024-12-27T23:28:54.5232Z","iopub.execute_input":"2024-12-27T23:28:54.523433Z","iopub.status.idle":"2024-12-27T23:28:55.742651Z","shell.execute_reply.started":"2024-12-27T23:28:54.523406Z","shell.execute_reply":"2024-12-27T23:28:55.74172Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Load Data","metadata":{}},{"cell_type":"code","source":"input_path = './input_df' if os.path.exists('./input_df') else '/content/'\nTRAINING = True\nfeature_names = [f\"feature_{i:02d}\" for i in range(79)] + [f\"responder_{idx}_lag_1\" for idx in range(9)]\nlabel_name = 'responder_6'\nweight_name = 'weight'\ntrain_name = os.path.join(\"./input_df/\", \"nn_input_df_with_lags.pickle\")\nvalid_name = os.path.join(\"./input_df/\", \"nn_valid_df_with_lags.pickle\")\nos.makedirs(os.path.dirname(train_name), exist_ok=True)\nif TRAINING and not os.path.exists(train_name):\n    df = pl.scan_parquet(f\"{input_path}/training.parquet\").collect().to_pandas()\n    valid = pl.scan_parquet(f\"{input_path}/validation.parquet\").collect().to_pandas()\n    df = pd.concat([df, valid]).reset_index(drop=True)# A trick to boost LB from 0.0045->0.005\n    with open(train_name, \"wb\") as w:\n        pickle.dump(df, w)\n    with open(valid_name, \"wb\") as w:\n        pickle.dump(valid, w)\nelif TRAINING:\n    with open(train_name, \"rb\") as r:\n        df = pickle.load(r)\n    with open(valid_name, \"rb\") as r:\n        valid = pickle.load(r)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-27T23:29:10.243376Z","iopub.execute_input":"2024-12-27T23:29:10.243722Z","execution_failed":"2024-12-27T23:30:16.235Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X_train = df[ feature_names ]\ny_train = df[ label_name ]\nw_train = df[ \"weight\" ]\nX_valid = valid[ feature_names ]\ny_valid = valid[ label_name ]\nw_valid = valid[ \"weight\" ]\n\nX_train.shape, y_train.shape, w_train.shape, X_valid.shape, y_valid.shape, w_valid.shape","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Training Configurations","metadata":{}},{"cell_type":"code","source":"import os\nimport warnings\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport pytorch_lightning as pl\nfrom pytorch_lightning import (LightningDataModule, LightningModule, Trainer)\nfrom pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint, Timer\nfrom pytorch_lightning.loggers import WandbLogger\nimport wandb\nimport pandas as pd\nimport numpy as np\nfrom sklearn.metrics import r2_score\nfrom sklearn.model_selection import train_test_split\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.model_selection._split import _BaseKFold, indexable, _num_samples\nfrom sklearn.utils.validation import _deprecate_positional_args\n\nclass custom_args():\n    def __init__(self):\n        self.usegpu = True\n        self.gpuid = 0\n        self.seed = 42\n        self.model = 'nn'\n        self.use_wandb = False\n        self.project = 'js-xs-nn-with-lags'\n        self.dname = \"./input_df/\"\n        self.loader_workers = 8\n        self.bs = 16384\n        self.lr = 1e-3\n        self.weight_decay = 5e-4\n        self.recon_weight = 0.05\n        self.dropouts = [0.1, 0.2]\n        self.n_hidden = [512, 512, 256]\n        self.patience = 25\n        self.max_epochs = 200\n        self.N_fold = 5\n\n\nmy_args = custom_args()","metadata":{"execution":{"iopub.status.busy":"2024-11-05T04:58:17.964754Z","iopub.execute_input":"2024-11-05T04:58:17.965474Z","iopub.status.idle":"2024-11-05T04:58:17.993759Z","shell.execute_reply.started":"2024-11-05T04:58:17.965427Z","shell.execute_reply":"2024-11-05T04:58:17.992851Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# PyTorch Data Module Definition","metadata":{}},{"cell_type":"code","source":"class PurgedGroupTimeSeriesSplit:\n    \"\"\"\n    A minimal example of a custom purged time-series split.\n    This is not official scikit-learn code; adjust as needed!\n    \"\"\"\n\n    def __init__(self, n_splits=5, group_gap=20):\n        self.n_splits = n_splits\n        self.group_gap = group_gap  # how many 'time units' to skip/purge before/after the val set\n\n    def split(self, X, groups):\n        \"\"\"\n        :param X: Can be your feature DataFrame or just an index array.\n        :param groups: The 'date_id' or time grouping array, sorted in ascending order ideally.\n        Y is ignored for simplicity.\n        \"\"\"\n        unique_groups = np.unique(groups)\n        n_samples = len(unique_groups)\n\n        fold_size = n_samples // self.n_splits\n\n        # Sequential chunk\n        for i in range(self.n_splits):\n            # The validation fold is chunk i\n            val_start = i * fold_size\n            val_end = val_start + fold_size\n            if i == self.n_splits - 1:\n                val_end = n_samples\n\n            # Indices for validation chunk in terms of group/time\n            val_groups = unique_groups[val_start:val_end]\n\n            # Purge region on either side\n            train_groups_left = unique_groups[: max(0, val_start - self.group_gap)]\n            train_groups_right = unique_groups[min(n_samples, val_end + self.group_gap):]\n\n            train_groups = np.concatenate([train_groups_left, train_groups_right])\n\n            # Now convert group arrays back to sample indices\n            train_idx = np.where(np.isin(groups, train_groups))[0]\n            val_idx   = np.where(np.isin(groups, val_groups))[0]\n\n            yield train_idx, val_idx","metadata":{"execution":{"iopub.status.busy":"2024-11-05T04:58:17.994912Z","iopub.execute_input":"2024-11-05T04:58:17.995267Z","iopub.status.idle":"2024-11-05T04:58:18.776454Z","shell.execute_reply.started":"2024-11-05T04:58:17.995223Z","shell.execute_reply":"2024-11-05T04:58:18.77564Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CustomDataset(Dataset):\n    def __init__(self, df, accelerator, scaler=None):\n        # If no scaler provided (training data), compute statistics\n        if scaler is None:\n            self.mean = df[feature_names].mean()\n            self.std = df[feature_names].std()\n            self.scaler = {'mean': self.mean, 'std': self.std}\n        else:\n            # Use provided scaler (for validation/test data)\n            self.mean = scaler['mean']\n            self.std = scaler['std']\n\n        # Normalize features\n        normalized_features = (df[feature_names] - self.mean) / (self.std + 1e-8)\n\n        self.features = torch.FloatTensor(normalized_features.values).to(accelerator)\n        self.labels = torch.FloatTensor(df[label_name].values).to(accelerator)\n        self.weights = torch.FloatTensor(df[weight_name].values).to(accelerator)\n\n    def __len__(self):\n        return len(self.labels)\n\n    def __getitem__(self, idx):\n        x = self.features[idx]\n        y = self.labels[idx]\n        w = self.weights[idx]\n        return x, y, w\n\n\nclass DataModule(LightningDataModule):\n    def __init__(self, train_df, batch_size, valid_df=None, accelerator='cpu', n_splits=5, group_gap=1):\n        super().__init__()\n        self.df = train_df\n        self.batch_size = batch_size\n        self.accelerator = accelerator\n        self.valid_df = valid_df\n        self.val_dataset = None\n        self.groups = self.df['date_id'].values\n        self.n_splits = n_splits\n        self.group_gap = group_gap\n        self.train_dataset = None\n\n    def setup(self, fold=0, N_fold=5, stage=None):\n        purged_splitter = PurgedGroupTimeSeriesSplit(n_splits=self.n_splits, group_gap=self.group_gap)\n        X_idx = np.arange(len(self.df))\n\n        train_idx, val_idx = None, None\n        for i, (tr, va) in enumerate(purged_splitter.split(X_idx, self.groups)):\n            if i == fold:\n                train_idx, val_idx = tr, va\n                break\n\n        if train_idx is not None and val_idx is not None:\n            df_train = self.df.iloc[train_idx]\n            df_val = self.df.iloc[val_idx]\n        else:\n            raise ValueError(f\"Fold {fold} is out of range for n_splits={self.n_splits}.\")\n\n        # Create train dataset first (no scaler provided, it will compute statistics)\n        self.train_dataset = CustomDataset(df_train, self.accelerator, scaler=None)\n\n        # Save scaler for inference\n        torch.save(self.train_dataset.scaler, 'scaler.pt')\n\n        # Use same scaler for validation\n        if self.valid_df is None:\n            self.val_dataset = CustomDataset(df_val, self.accelerator, scaler=self.train_dataset.scaler)\n        else:\n            self.val_dataset = CustomDataset(self.valid_df, self.accelerator, scaler=self.train_dataset.scaler)\n\n    def train_dataloader(self, n_workers=0):\n        return DataLoader(\n            self.train_dataset,\n            batch_size=self.batch_size,\n            shuffle=False,  # Keep temporal order\n            num_workers=n_workers\n        )\n\n    def val_dataloader(self, n_workers=0):\n        return DataLoader(\n            self.val_dataset,\n            batch_size=self.batch_size,\n            shuffle=False,\n            num_workers=n_workers\n        )\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# NN Model Definition","metadata":{}},{"cell_type":"code","source":"class SupervisedAE(LightningModule):\n    def __init__(self, input_dim, hidden_dims, dropouts, lr, weight_decay, recon_weight, noise_std=0.1):\n        super().__init__()\n        self.save_hyperparameters()\n        self.noise_std = noise_std\n        self.lr = lr\n        self.weight_decay = weight_decay\n        self.recon_weight = recon_weight\n\n        # Encoder with residual connections\n        self.encoder_layers = nn.ModuleList()\n        in_dim = input_dim\n        for hd in hidden_dims:\n            self.encoder_layers.append(nn.Sequential(\n                nn.Linear(in_dim, hd),\n                nn.SiLU(),\n                nn.Dropout(dropouts[1])\n            ))\n            if in_dim == hd:  # Add residual connection if dimensions match\n                self.encoder_layers[-1].residual = True\n            else:\n                self.encoder_layers[-1].residual = False\n            in_dim = hd\n\n        # Decoder with similar architecture\n        self.decoder_layers = nn.ModuleList()\n        rev_dims = list(reversed(hidden_dims))\n        in_dim = rev_dims[0]\n        for hd in rev_dims[1:]:\n            self.decoder_layers.append(nn.Sequential(\n                nn.Linear(in_dim, hd),\n                nn.SiLU(),\n                nn.Dropout(dropouts[0])\n            ))\n            in_dim = hd\n        self.decoder_layers.append(nn.Linear(in_dim, input_dim))\n\n        # Single prediction head for responder_6\n        self.head = nn.Sequential(\n            nn.Linear(hidden_dims[-1], hidden_dims[-1] // 2),\n            nn.SiLU(),\n            nn.Dropout(dropouts[0]),\n            nn.Linear(hidden_dims[-1] // 2, 1)\n        )\n\n        self.validation_step_outputs = []\n\n    def forward(self, x):\n        # Add noise during training\n        if self.training:\n            x = x + torch.randn_like(x) * self.noise_std\n\n        # Encoder with residual connections\n        z = x\n        for layer in self.encoder_layers:\n            if layer.residual:\n                z = z + layer(z)\n            else:\n                z = layer(z)\n\n        # Decoder\n        x_hat = z\n        for layer in self.decoder_layers:\n            x_hat = layer(x_hat)\n\n        # Prediction for responder_6\n        y_lin = self.head(z)\n        y_hat = 5.0 * torch.tanh(y_lin).squeeze(-1) \n\n        return x_hat, y_hat\n\n    def training_step(self, batch, batch_idx):\n        x, y, w = batch\n        x_hat, y_hat = self(x)\n\n        # Reconstruction loss\n        recon_loss = F.mse_loss(x_hat, x, reduction=\"none\").mean(dim=1)\n\n        # Supervised loss for responder_6\n        sup_loss = F.mse_loss(y_hat, y, reduction=\"none\")\n\n        # Combine losses with weights\n        total_loss = (self.recon_weight * recon_loss + sup_loss) * w\n        loss = total_loss.mean()\n\n        # Log metrics\n        self.log(\"train_loss\", loss)\n        self.log(\"train_recon_loss\", recon_loss.mean())\n        self.log(\"train_sup_loss\", sup_loss.mean())\n\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        x, y, w = batch\n        x_hat, y_hat = self(x)\n\n        # Reconstruction loss\n        recon_loss = F.mse_loss(x_hat, x, reduction=\"none\").mean(dim=1)\n\n        # Supervised loss for responder_6\n        sup_loss = F.mse_loss(y_hat, y, reduction=\"none\")\n\n        # Combine losses with weights\n        total_loss = (self.recon_weight * recon_loss + sup_loss) * w\n        val_loss = total_loss.mean()\n\n        # Log validation metrics\n        self.log(\"val_loss\", val_loss, prog_bar=True, on_step=False, on_epoch=True)\n        self.log(\"val_recon_loss\", recon_loss.mean(), on_step=False, on_epoch=True)\n        self.log(\"val_sup_loss\", sup_loss.mean(), on_step=False, on_epoch=True)\n\n        # Store predictions for R^2 calculation\n        self.validation_step_outputs.append({\n            'y_hat': y_hat.detach(),\n            'y_true': y.detach(),\n            'weights': w.detach()\n        })\n\n        return val_loss\n\n    def on_validation_epoch_end(self):\n        \"\"\"Calculate validation weighted R^2 at the end of the epoch.\"\"\"\n        # Concatenate all batches\n        y_hat = torch.cat([x['y_hat'] for x in self.validation_step_outputs]).cpu().numpy()\n        y_true = torch.cat([x['y_true'] for x in self.validation_step_outputs]).cpu().numpy()\n        weights = torch.cat([x['weights'] for x in self.validation_step_outputs]).cpu().numpy()\n\n        # Calculate weighted R^2\n        numerator = np.average((y_hat - y_true)**2, weights=weights)\n        denominator = np.average(y_true**2, weights=weights) + 1e-38\n        val_r2 = 1.0 - numerator / denominator\n\n        # Log metrics\n        self.log(\"val_r_square\", val_r2, prog_bar=True)\n\n        # Clear stored outputs\n        self.validation_step_outputs.clear()\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)\n        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n            optimizer,\n            mode='min',\n            factor=0.5,\n            patience=5,\n            verbose=True\n        )\n        return {\n            'optimizer': optimizer,\n            'lr_scheduler': {\n                'scheduler': scheduler,\n                'monitor': 'val_loss',\n            }\n        }\n\n    def on_train_epoch_end(self):\n        if self.trainer.sanity_checking:\n            return\n        epoch = self.trainer.current_epoch\n        metrics = {k: v.item() if isinstance(v, torch.Tensor) else v\n                  for k, v in self.trainer.logged_metrics.items()}\n        formatted_metrics = {k: f\"{v:.5f}\" for k, v in metrics.items()}\n        print(f\"Epoch {epoch}: {formatted_metrics}\")","metadata":{"execution":{"iopub.status.busy":"2024-11-05T04:58:18.778669Z","iopub.execute_input":"2024-11-05T04:58:18.778953Z","iopub.status.idle":"2024-11-05T04:58:18.796672Z","shell.execute_reply.started":"2024-11-05T04:58:18.778923Z","shell.execute_reply":"2024-11-05T04:58:18.795818Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Create PyTorch Data Module","metadata":{}},{"cell_type":"code","source":"args = my_args\n\n# checking device\ndevice = torch.device(f'cuda:{args.gpuid}' if torch.cuda.is_available() and args.usegpu else 'cpu')\naccelerator = 'gpu' if torch.cuda.is_available() and args.usegpu else 'cpu'\nloader_device = 'cpu'\n\n\n# Initialize Data Module\n\ndf[feature_names] = df[feature_names].fillna(method = 'ffill').fillna(0)\nvalid[feature_names] = valid[feature_names].fillna(method = 'ffill').fillna(0)\ndata_module = DataModule(df, batch_size=args.bs, valid_df=valid, accelerator=loader_device)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Create Model and Training","metadata":{}},{"cell_type":"code","source":"import gc\ndel df\ngc.collect()\npl.seed_everything(args.seed)\nfor fold in range(args.N_fold):\n    data_module.setup(fold, args.N_fold)\n    # Obtain input dimension\n    input_dim = data_module.train_dataset.features.shape[1]\n    # Initialize Model\n    model = SupervisedAE(\n        input_dim=input_dim,\n        hidden_dims=args.n_hidden,\n        dropouts=args.dropouts,\n        lr=args.lr,\n        weight_decay=args.weight_decay,\n        recon_weight=0.05\n    )\n    # Initialize Logger\n    if args.use_wandb:\n        wandb_run = wandb.init(project=args.project, config=vars(args), reinit=True)\n        logger = WandbLogger(experiment=wandb_run)\n    else:\n        logger = None\n    # Initialize Callbacks\n    early_stopping = EarlyStopping('val_loss', patience=args.patience, mode='min', verbose=False)\n    checkpoint_callback = ModelCheckpoint(monitor='val_loss', mode='min', save_top_k=1, verbose=False, filename=f\"./models/nn_{fold}.model\", dirpath='/content/drive/MyDrive/Colab Notebooks/Jane Street/models')\n    timer = Timer()\n    # Initialize Trainer\n    trainer = Trainer(\n        max_epochs=args.max_epochs,\n        accelerator=accelerator,\n        devices=[args.gpuid] if args.usegpu else None,\n        logger=logger,\n        callbacks=[early_stopping, checkpoint_callback, timer],\n        enable_progress_bar=True\n    )\n    # Start Training\n    trainer.fit(model, data_module.train_dataloader(args.loader_workers), data_module.val_dataloader(args.loader_workers))\n    # You can find trained best model in your local path\n    print(f'Fold-{fold} Training completed in {timer.time_elapsed(\"train\"):.2f}s')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-05T04:58:18.797992Z","iopub.execute_input":"2024-11-05T04:58:18.798262Z","iopub.status.idle":"2024-11-05T04:58:18.818565Z","shell.execute_reply.started":"2024-11-05T04:58:18.798232Z","shell.execute_reply":"2024-11-05T04:58:18.817654Z"}},"outputs":[],"execution_count":null}]}