{"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): **reference 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":"markdown","source":"### This notebook is based on https://www.kaggle.com/code/voix97/jane-street-rmf-training-nn. \n### It has been rewritten using Polars instead of Pandas. Additionally, the training is done using PyTorch alone instead of PyTorch Lightning. Although training takes time, it can also run on Kaggle Notebooks.","metadata":{}},{"cell_type":"markdown","source":"# Training Neural Networks (MLP)","metadata":{}},{"cell_type":"markdown","source":"# Library","metadata":{}},{"cell_type":"code","source":"\n# ====================================================\n# Library\n# ====================================================\nimport os\nimport gc\nimport time\nimport random\nimport polars as pl\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nfrom tqdm import tqdm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-29T02:18:57.693154Z","iopub.execute_input":"2024-12-29T02:18:57.693405Z","iopub.status.idle":"2024-12-29T02:19:01.049965Z","shell.execute_reply.started":"2024-12-29T02:18:57.693372Z","shell.execute_reply":"2024-12-29T02:19:01.049037Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Config","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# CFG\n# ====================================================\nclass CFG:\n    seed = 42\n    use_wandb = False\n    apex=True\n    exp = 'js-exp001'\n    model = 'NN'\n    model_path = '.'\n    input_path = '/kaggle/input/js24-preprocessing-create-lags'\n    loader_workers = 4\n    batch_size = 8192\n    lr = 1e-3\n    weight_decay = 5e-4\n    max_grad_norm=1000\n    dropouts = [0.1, 0.1]\n    n_hidden = [512, 512, 256]\n    patience = 25\n    # epochs = 2000\n    epochs = 15 #demo\n    n_fold=5\n    # trn_fold = [0,1,2,3,4]\n    trn_fold=[0] #demo\n\n    feature_names = [f\"feature_{i:02d}\" for i in range(79)] + [f\"responder_{idx}_lag_1\" for idx in range(9)]\n    label_name = 'responder_6'\n    weight_name = 'weight'\n\ndevice = torch.device(f'cuda:0' if torch.cuda.is_available() else 'cpu')\n\n# ====================================================\n# wandb\n# ====================================================\nif CFG.use_wandb:\n    \n    import wandb\n\n    try:\n        from kaggle_secrets import UserSecretsClient\n        user_secrets = UserSecretsClient()\n        secret_value_0 = user_secrets.get_secret(\"wandb\")\n        wandb.login(key=secret_value_0)\n        anony = None\n    except:\n        anony = \"must\"\n        print('If you want to use your W&B account, go to Add-ons -> Secrets and provide your W&B access token. Use the Label name as wandb_api. \\nGet your W&B access token from here: https://wandb.ai/authorize')\n\n\n    def class2dict(f):\n        return dict((name, getattr(f, name)) for name in dir(f) if not name.startswith('__'))\n\n    run = wandb.init(project='jane-street-24', \n                     name=CFG.model,\n                     config=class2dict(CFG),\n                     group=CFG.model,\n                     job_type=\"train\",\n                     anonymous=anony)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-29T02:19:01.051836Z","iopub.execute_input":"2024-12-29T02:19:01.052216Z","iopub.status.idle":"2024-12-29T02:19:01.1141Z","shell.execute_reply.started":"2024-12-29T02:19:01.05219Z","shell.execute_reply":"2024-12-29T02:19:01.112854Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Utils","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Utils\n# ====================================================\n# Custom R2 metric for validation\ndef r2_val(y_true, y_pred, sample_weight):\n    r2 = 1 - np.average((y_pred - y_true) ** 2, weights=sample_weight) / (np.average((y_true) ** 2, weights=sample_weight) + 1e-38)\n    return r2\n    \ndef seed_everything(seed=42):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    \nseed_everything(seed=CFG.seed)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-29T02:19:01.115257Z","iopub.execute_input":"2024-12-29T02:19:01.115519Z","iopub.status.idle":"2024-12-29T02:19:01.133556Z","shell.execute_reply.started":"2024-12-29T02:19:01.11547Z","shell.execute_reply":"2024-12-29T02:19:01.132839Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Data Loading","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Data Loading\n# ====================================================\ndf = pl.scan_parquet(f\"{CFG.input_path}/training.parquet\").collect()\nvalid = pl.scan_parquet(f\"{CFG.input_path}/validation.parquet\").collect()\n# Comment out the trick to make sure it works. =======\n# df = pl.concat([df, valid]).with_row_count().select([pl.all().exclude(\"row_nr\")]) # A trick to boost LB from 0.0045->0.005\n\ndf = df.with_columns([pl.col(CFG.feature_names).fill_null(strategy=\"forward\").fill_null(0)])\nvalid = valid.with_columns([pl.col(CFG.feature_names).fill_null(strategy=\"forward\").fill_null(0)])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-29T02:19:01.134581Z","iopub.execute_input":"2024-12-29T02:19:01.134794Z","iopub.status.idle":"2024-12-29T02:19:30.709322Z","shell.execute_reply.started":"2024-12-29T02:19:01.134771Z","shell.execute_reply":"2024-12-29T02:19:30.708635Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Dataset\n# ====================================================\nclass CustomDataset(Dataset):\n    def __init__(self, df, config):\n        self.features = torch.tensor(df.select(config.feature_names).to_numpy(), dtype=torch.float32)\n        self.labels = torch.tensor(df.select(config.label_name).to_numpy(), dtype=torch.float32).flatten()\n        self.weights = torch.tensor(df.select(config.weight_name).to_numpy(), dtype=torch.float32).flatten()\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    def __len__(self):\n        return len(self.labels)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-29T02:19:30.710418Z","iopub.execute_input":"2024-12-29T02:19:30.710775Z","iopub.status.idle":"2024-12-29T02:19:30.716996Z","shell.execute_reply.started":"2024-12-29T02:19:30.710738Z","shell.execute_reply":"2024-12-29T02:19:30.716098Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Model\n# ====================================================\nclass NN(nn.Module):\n    def __init__(self, input_dim, hidden_dims, dropouts):\n        super(NN, self).__init__()\n        layers = []\n        in_dim = input_dim\n        for i, hidden_dim in enumerate(hidden_dims):\n            layers.append(nn.BatchNorm1d(in_dim))\n            if i > 0:\n                layers.append(nn.SiLU())\n            if i < len(dropouts):\n                layers.append(nn.Dropout(dropouts[i]))\n            layers.append(nn.Linear(in_dim, hidden_dim))\n            in_dim = hidden_dim\n        layers.append(nn.Linear(in_dim, 1))\n        layers.append(nn.Tanh())\n        self.model = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return 5 * self.model(x).squeeze(-1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-29T02:19:30.718122Z","iopub.execute_input":"2024-12-29T02:19:30.718372Z","iopub.status.idle":"2024-12-29T02:19:30.728267Z","shell.execute_reply.started":"2024-12-29T02:19:30.718346Z","shell.execute_reply":"2024-12-29T02:19:30.727372Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Train loop","metadata":{}},{"cell_type":"code","source":"def train_fn(model, train_loader, optimizer, criterion, config, device):\n    model.train()\n    train_loss = 0\n    scaler = torch.amp.GradScaler()\n\n    for batch in tqdm(train_loader, desc=\"Training\", leave=False):\n        x, y, w = batch\n        x, y, w = x.to(device), y.to(device), w.to(device) \n        optimizer.zero_grad()\n        with torch.amp.autocast(device_type='cuda', dtype=torch.float16, enabled=config.apex):\n            y_hat = model(x)\n            loss = criterion(y_hat, y) *  w\n        loss = loss.mean()\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm)\n        scaler.update()\n        train_loss += loss.item() * x.size(0)\n    train_loss /= len(train_loader.dataset) \n    return train_loss\n\n\ndef validate_fn(model, val_loader, criterion, config, device):\n    model.eval()\n    val_loss = 0\n    y_true, y_pred, weights = [], [], []\n    for batch in tqdm(val_loader, desc=\"Validation\", leave=False):\n        x, y, w = batch\n        x, y, w = x.to(device), y.to(device), w.to(device)\n        with torch.no_grad():\n            y_hat = model(x)\n            loss = criterion(y_hat, y) *  w\n        loss = loss.mean()\n        val_loss += loss.item() * x.size(0)\n        y_true.append(y.cpu().numpy())\n        y_pred.append(y_hat.cpu().numpy())\n        weights.append(w.cpu().numpy())\n    val_loss /= len(val_loader.dataset)\n    y_true = np.concatenate(y_true)\n    y_pred = np.concatenate(y_pred)\n    weights = np.concatenate(weights)\n    val_r2 = r2_val(y_true, y_pred, weights)\n    return val_loss, val_r2\n    \n# ====================================================\n# train loop\n# ====================================================\ndef train_loop(df, fold, config):\n    print(f\"========== fold: {fold} training ==========\")\n    # ====================================================\n    # loader\n    # ====================================================\n    dates = df.select('date_id').unique().to_series().to_list()\n    selected_dates = [date for ii, date in enumerate(dates) if ii % config.n_fold != fold]\n    df_train = df.filter(pl.col('date_id').is_in(selected_dates))\n    train_dataset = CustomDataset(df_train, config)\n    valid_dataset = CustomDataset(valid, config)\n    train_loader  = DataLoader(train_dataset,\n                               batch_size=config.batch_size, \n                               shuffle=True,\n                               num_workers=config.loader_workers, \n                               pin_memory=True)\n    valid_loader = DataLoader(valid_dataset,\n                               batch_size=config.batch_size, \n                               shuffle=False, \n                               num_workers=config.loader_workers,\n                               pin_memory=True)\n\n    # ====================================================\n    # model & optimizer\n    # ====================================================\n    best_score = -np.inf\n    patience_counter = 0\n    \n    input_dim = len(config.feature_names)\n    model = NN(input_dim=input_dim, hidden_dims=config.n_hidden, dropouts=config.dropouts)\n    # model = nn.DataParallel(model)\n    model.to(device)\n    optimizer = torch.optim.Adam(model.parameters(), lr=config.lr, weight_decay=config.weight_decay)\n    num_train_steps = int(len(df_train) / config.batch_size * config.epochs)\n    scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5)\n    criterion = nn.MSELoss(reduction='none')\n\n    for epoch in range(config.epochs):\n        start_time = time.time()\n        print(f\"Epoch {epoch+1}/{config.epochs}\")\n        train_loss = train_fn(model, train_loader, optimizer, criterion, config, device)\n        val_loss, score = validate_fn(model, valid_loader, criterion, config, device)\n        scheduler.step(val_loss)\n\n        print(f\"Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, lr: {scheduler.get_last_lr()[0]:.4f}, Val R2: {score:.4f}\")\n\n        if config.use_wandb:\n            wandb.log({f\"[fold{fold}] epoch\": epoch+1, \n                       f\"[fold{fold}] avg_train_loss\": train_loss, \n                       f\"[fold{fold}] avg_val_loss\": val_loss,\n                       f\"[fold{fold}] score\": score,\n                       f\"[fold{fold}] lr\": scheduler.get_last_lr()[0],\n                       })\n        if score > best_score:\n            best_score = score\n            patience_counter = 0\n            torch.save({\n                        'epoch': epoch,\n                        'model_state_dict': model.state_dict(),\n                        'best_score': best_score,\n                    }, f\"{model_path}/model_checkpoint_fold{fold}.pth\")\n            print(f\"Save Best Score: {score:.4f} \")\n        else:\n            patience_counter += 1\n            if patience_counter >= config.patience:\n                print(f\"Early stopping triggered at epoch {epoch+1}\")\n                break\n    del train_dataset, valid_dataset, train_loader, valid_loader, model, optimizer, scheduler\n    gc.collect()\n    torch.cuda.empty_cache()\n\nif __name__ == \"__main__\":\n\n    config = CFG()\n    model_path = os.path.join(config.model_path, config.exp)\n    if not os.path.exists(model_path):\n        os.makedirs(model_path)\n    for fold in config.trn_fold:\n        print(f\"Starting fold {fold + 1}/{config.n_fold}\")\n        train_loop(df, fold, config)\n    if config.use_wandb:\n        wandb.finish()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-29T02:19:30.731166Z","iopub.execute_input":"2024-12-29T02:19:30.731588Z","execution_failed":"2024-12-29T02:20:22.71Z"}},"outputs":[],"execution_count":null}]}