{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":53666,"databundleVersionId":6589269,"sourceType":"competition"},{"sourceId":6684565,"sourceType":"datasetVersion","datasetId":3714177}],"dockerImageVersionId":30786,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport polars as pl\nfrom tqdm.auto import tqdm\nfrom matplotlib import pyplot as plt\nfrom sklearn.model_selection import GroupKFold\nfrom sklearn.metrics import mean_squared_error\nfrom scipy.signal import find_peaks\nimport gc\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\nfrom torchinfo import summary\n\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:37:16.09452Z","iopub.execute_input":"2024-12-01T16:37:16.094955Z","iopub.status.idle":"2024-12-01T16:37:23.075456Z","shell.execute_reply.started":"2024-12-01T16:37:16.09492Z","shell.execute_reply":"2024-12-01T16:37:23.074197Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_series = pl.read_parquet(\"/kaggle/input/child-mind-institute-detect-sleep-states/train_series.parquet\")\ndf_events = pl.read_csv(\"/kaggle/input/child-mind-institute-detect-sleep-states/train_events.csv\")\ndf_events = df_events.with_columns(pl.col(\"event\").replace_strict({\"wakeup\": 1.0, \"onset\": -1.0}, return_dtype=pl.Float32))\ndf_groupby = df_series.group_by(\"series_id\", maintain_order=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:37:27.442482Z","iopub.execute_input":"2024-12-01T16:37:27.443106Z","iopub.status.idle":"2024-12-01T16:37:40.639085Z","shell.execute_reply.started":"2024-12-01T16:37:27.443065Z","shell.execute_reply":"2024-12-01T16:37:40.637796Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_series","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:37:40.641054Z","iopub.execute_input":"2024-12-01T16:37:40.641444Z","iopub.status.idle":"2024-12-01T16:37:40.658625Z","shell.execute_reply.started":"2024-12-01T16:37:40.641406Z","shell.execute_reply":"2024-12-01T16:37:40.657356Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_events","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:37:49.394814Z","iopub.execute_input":"2024-12-01T16:37:49.395249Z","iopub.status.idle":"2024-12-01T16:37:49.403925Z","shell.execute_reply.started":"2024-12-01T16:37:49.395209Z","shell.execute_reply":"2024-12-01T16:37:49.402565Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_groupby","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:38:07.680081Z","iopub.execute_input":"2024-12-01T16:38:07.680555Z","iopub.status.idle":"2024-12-01T16:38:07.689133Z","shell.execute_reply.started":"2024-12-01T16:38:07.680514Z","shell.execute_reply":"2024-12-01T16:38:07.687569Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dict_valid_ratio = dict()\nlist_feature_array = []\nlist_df_1min = []\nif df_events.schema[\"series_id\"] == pl.List:\n    df_events = df_events.explode(\"series_id\")\n\n# Ensure series_id is a scalar and compatible\ndf_events = df_events.with_columns(pl.col(\"series_id\").cast(str))\nfor series_id, df in tqdm(df_groupby, total=df_series.get_column(\"series_id\").n_unique()):\n    if isinstance(series_id, tuple):\n        series_id = series_id[0]  # Extract the relevant part of the tuple\n    series_id = str(series_id)\n    df = (\n        df.join(\n            df_events.filter(pl.col(\"series_id\") == series_id).select(\"timestamp\", \"event\"),\n            on=\"timestamp\", how=\"left\"\n        ).with_columns(\n            pl.col(\"timestamp\").str.to_datetime(),\n            pl.col(\"event\").fill_null(0.0),\n        ).with_columns(\n            pl.col(\"timestamp\").dt.date().cast(str).alias(\"date\"),\n            pl.col(\"timestamp\").dt.time().cast(str).alias(\"time\"),\n        )\n    ).to_pandas()\n    \n    #  Handle duplicate data and calculate validity flag\n    df[\"timestamp\"] = df[\"timestamp\"].dt.tz_localize(None)\n    dup_count = df.groupby([\"anglez\", \"enmo\", \"time\"])[\"step\"].transform(\"count\")\n    df[\"valid_flag\"] = (dup_count == 1).astype(\"float32\")\n    dict_valid_ratio[series_id] = df[\"valid_flag\"].mean()\n    \n    #     Calculate rolling features and pivot the data\n    list_feature_array_tmp = []\n    df[\"log_anglez_std\"] = np.log(df[\"anglez\"].rolling(12, min_periods=1, center=True).std() + 1).astype(\"float32\")\n    df[\"log_enmo\"] = np.log(df[\"enmo\"] + 0.01).astype(\"float32\")\n    for feature in [\"log_anglez_std\", \"log_enmo\", \"valid_flag\"]:\n        df_pivot = df.pivot(index=[\"series_id\", \"date\"], columns=\"time\", values=feature)\n        feature_array = df_pivot.fillna(0).values\n        feature_array_1day_bedore = df_pivot.shift(1).fillna(0).values\n        feature_array_1day_after = df_pivot.shift(-1).fillna(0).values\n        feature_array = np.concatenate([feature_array_1day_bedore[:, -180*12:], feature_array, feature_array_1day_after[:, :180*12]], axis=1)\n        list_feature_array_tmp.append(feature_array)\n    list_feature_array.append(np.stack(list_feature_array_tmp, axis=1))\n    \n    #     Resample data to 1 minute intervals\n    dict_agg = {\"series_id\": \"first\", \"date\": \"first\", \"time\": \"first\", \"step\": \"mean\", \"event\": \"sum\", \"valid_flag\": \"max\"}\n    df_1min = df.resample(\"1min\", on=\"timestamp\").agg(dict_agg).reset_index()\n    df_1min[\"step\"] = df_1min[\"step\"].astype(\"int32\")\n    \n    #     Compute target values for events\n    values_event = df_1min[\"event\"].values\n    values_target = values_event.copy()\n    \n    #     window of 30 minutes (or steps) in both backward and forward directions\n    for j in range(30):\n        weight = np.exp(-j/2.8)\n        values_target[:-(j+1)] += weight*values_event[(j+1):]  # shift backward\n        if j > 0:\n            values_target[j:] += weight*values_event[:-j]  # shift forward\n    \n    df_1min[\"target\"] = values_target\n    list_df_1min.append(df_1min)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:38:17.389961Z","iopub.execute_input":"2024-12-01T16:38:17.390427Z","iopub.status.idle":"2024-12-01T16:44:51.014333Z","shell.execute_reply.started":"2024-12-01T16:38:17.390388Z","shell.execute_reply":"2024-12-01T16:44:51.012557Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(df_events.schema)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:22:22.496414Z","iopub.execute_input":"2024-12-01T16:22:22.496841Z","iopub.status.idle":"2024-12-01T16:22:22.503334Z","shell.execute_reply.started":"2024-12-01T16:22:22.496803Z","shell.execute_reply":"2024-12-01T16:22:22.502014Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"del df_series, df_groupby\ngc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:45:10.872579Z","iopub.execute_input":"2024-12-01T16:45:10.873055Z","iopub.status.idle":"2024-12-01T16:45:11.073334Z","shell.execute_reply.started":"2024-12-01T16:45:10.873005Z","shell.execute_reply":"2024-12-01T16:45:11.072107Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Normalizing the Array\n# This operation standardizes each feature's values across all timesteps and samples to be between 0 and 1\n\nX = np.concatenate(list_feature_array)\nX = (X - X.min(axis=(0, 2), keepdims=True)) / (X.max(axis=(0, 2), keepdims=True) - X.min(axis=(0, 2), keepdims=True))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:45:44.697081Z","iopub.execute_input":"2024-12-01T16:45:44.697516Z","iopub.status.idle":"2024-12-01T16:45:47.394066Z","shell.execute_reply.started":"2024-12-01T16:45:44.697478Z","shell.execute_reply":"2024-12-01T16:45:47.392388Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:45:57.608978Z","iopub.execute_input":"2024-12-01T16:45:57.609431Z","iopub.status.idle":"2024-12-01T16:45:57.617029Z","shell.execute_reply.started":"2024-12-01T16:45:57.609391Z","shell.execute_reply":"2024-12-01T16:45:57.615774Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_1min = pd.concat(list_df_1min)\ndf_1min","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:46:05.686238Z","iopub.execute_input":"2024-12-01T16:46:05.686827Z","iopub.status.idle":"2024-12-01T16:46:06.419418Z","shell.execute_reply.started":"2024-12-01T16:46:05.686774Z","shell.execute_reply":"2024-12-01T16:46:06.418199Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"del list_df_1min\ndel list_feature_array\ngc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:46:16.762879Z","iopub.execute_input":"2024-12-01T16:46:16.763347Z","iopub.status.idle":"2024-12-01T16:46:17.024432Z","shell.execute_reply.started":"2024-12-01T16:46:16.763299Z","shell.execute_reply":"2024-12-01T16:46:17.022966Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# reshape df_1min from a long format to a wide format \n\ndf_y = df_1min.pivot(index=[\"series_id\", \"date\"], columns=\"time\", values=\"target\").fillna(0)\ndf_mask = df_1min.pivot(index=[\"series_id\", \"date\"], columns=\"time\", values=\"valid_flag\").fillna(0)\ndf_mask.to_parquet(\"df_mask.parquet\")\n\ndf_y","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:46:25.75303Z","iopub.execute_input":"2024-12-01T16:46:25.753474Z","iopub.status.idle":"2024-12-01T16:46:33.77639Z","shell.execute_reply.started":"2024-12-01T16:46:25.753434Z","shell.execute_reply":"2024-12-01T16:46:33.775144Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MyDataset(Dataset):\n    def __init__(self, X, Y, flag):\n        self.X = torch.FloatTensor(X)\n        if Y is not None:\n            self.Y = torch.FloatTensor(Y)\n        self.flag = torch.FloatTensor(flag)\n    \n    def __len__(self):\n        return self.X.shape[0]\n    \n    def __getitem__(self, idx):\n        if \"Y\" in dir(self):\n            return (self.X[idx], self.Y[idx], self.flag[idx])\n        else:\n            return (self.X[idx], torch.Tensor(), self.flag[idx])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:46:49.575921Z","iopub.execute_input":"2024-12-01T16:46:49.576347Z","iopub.status.idle":"2024-12-01T16:46:49.583499Z","shell.execute_reply.started":"2024-12-01T16:46:49.576308Z","shell.execute_reply":"2024-12-01T16:46:49.582213Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 1D convolutional layer,\n# batch normalization, \n# a ReLU activation function\n\nclass ConvBNReLU(nn.Module):\n    def __init__(self, in_channels, out_channels, kernel_size, stride, groups=1):\n        super().__init__()\n        \n        if stride == 1:\n            padding = \"same\"\n        else:\n            padding = (kernel_size - stride) // 2\n        self.layers = nn.Sequential(\n            nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=groups),\n            nn.BatchNorm1d(out_channels),\n            nn.ReLU()\n        )\n    \n    def forward(self, x):\n        x_out = self.layers(x)\n        return x_out","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:47:32.594884Z","iopub.execute_input":"2024-12-01T16:47:32.595406Z","iopub.status.idle":"2024-12-01T16:47:32.603405Z","shell.execute_reply.started":"2024-12-01T16:47:32.595365Z","shell.execute_reply":"2024-12-01T16:47:32.602044Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Squeeze-and-Excitation (SE) block\n\nclass SEBlock(nn.Module):\n    def __init__(self, n_channels, se_ratio):\n        super().__init__()\n        \n        self.layers = nn.Sequential(\n            nn.AdaptiveAvgPool1d(output_size=1),  #  Global Average Pooling\n            nn.Conv1d(n_channels, n_channels//se_ratio, kernel_size=1),\n            nn.ReLU(),\n            nn.Conv1d(n_channels//se_ratio, n_channels, kernel_size=1),\n            nn.Sigmoid()\n        )\n\n    def forward(self, x):\n        x_out = torch.mul(x, self.layers(x))\n        return x_out\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:47:44.823588Z","iopub.execute_input":"2024-12-01T16:47:44.824015Z","iopub.status.idle":"2024-12-01T16:47:44.833408Z","shell.execute_reply.started":"2024-12-01T16:47:44.823978Z","shell.execute_reply":"2024-12-01T16:47:44.831862Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# residual block\n\nclass ResBlock(nn.Module):\n    def __init__(self, n_channels, kernel_size, se_ratio):\n        super().__init__()\n        \n        self.layers = nn.Sequential(\n            ConvBNReLU(n_channels, n_channels, kernel_size, stride=1),\n            ConvBNReLU(n_channels, n_channels, kernel_size, stride=1),\n            SEBlock(n_channels, se_ratio)\n        )\n    \n    def forward(self, x):\n        x_re = self.layers(x)\n        x_out = x + x_re\n        return x_out\n    ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:47:52.499484Z","iopub.execute_input":"2024-12-01T16:47:52.499948Z","iopub.status.idle":"2024-12-01T16:47:52.507558Z","shell.execute_reply.started":"2024-12-01T16:47:52.49991Z","shell.execute_reply":"2024-12-01T16:47:52.505906Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# U-Net model adapted for 1-dimensional data\n\nclass UNet1d(nn.Module):\n    def __init__(self, input_channels, initial_channels, initial_kernel_size,\n                 down_channels, down_kernel_size, down_stride, res_depth, res_kernel_size, se_ratio, out_kernel_size):\n        super().__init__()\n        self.down_kernel_size = down_kernel_size\n        self.down_stride = down_stride\n        \n        self.initial_layers = ConvBNReLU(input_channels, initial_channels, initial_kernel_size, stride=1, groups=input_channels)\n        \n        self.down_layers = nn.ModuleList()\n        for i in range(len(down_channels)):\n            if i == 0:\n                in_channels = initial_channels\n            else:\n                in_channels = down_channels[i-1] + input_channels\n            out_channels = down_channels[i]\n            kernel_size = down_kernel_size[i]\n            stride = down_stride[i]\n            \n            block = []\n            block.append(ConvBNReLU(in_channels, out_channels, kernel_size, stride))\n            for j in range(res_depth):\n                block.append(ResBlock(out_channels, res_kernel_size, se_ratio))\n            self.down_layers.append(nn.Sequential(*block))\n        \n        self.up_layers = nn.ModuleList()\n        for i in range(len(down_channels)-1, 0, -1):\n            in_channels = out_channels + down_channels[i]\n            out_channels = down_channels[i]\n            kernel_size = down_kernel_size[i]\n            self.up_layers.append(ConvBNReLU(in_channels, out_channels, kernel_size, stride=1))\n        \n        self.out_layers = nn.Conv1d(down_channels[1], 1, out_kernel_size, padding=\"same\")\n    \n    def forward(self, x):\n        outs = []\n        x_avg = x\n        x = self.initial_layers(x)\n        \n        for i in range(len(self.down_layers)):\n            x_out = self.down_layers[i](x)\n            if i == len(self.down_layers) - 1:\n                x = x_out\n            else:\n                outs.append(x_out)\n                kernel_size = self.down_kernel_size[i]\n                stride = self.down_stride[i]\n                padding = (kernel_size - stride) // 2\n                x_avg = F.avg_pool1d(x_avg, kernel_size, stride, padding)\n                x = torch.cat([x_out, x_avg], dim=1)\n        \n        for i in range(len(self.up_layers)):\n            scale_factor = self.down_stride[-i-1]\n            x = F.interpolate(x, scale_factor=scale_factor, mode=\"linear\")\n            x = torch.cat([x, outs[-i-1]], dim=1)\n            x = self.up_layers[i](x)\n        \n        x_out = self.out_layers(x)\n        x_out = x_out[:, 0, 180:-180]\n        \n        return x_out","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:48:07.405506Z","iopub.execute_input":"2024-12-01T16:48:07.405962Z","iopub.status.idle":"2024-12-01T16:48:07.420989Z","shell.execute_reply.started":"2024-12-01T16:48:07.405922Z","shell.execute_reply":"2024-12-01T16:48:07.419378Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train(model, data_loader, optimizer, criterion, device):\n    model.train()\n    for batch in data_loader:\n        X = batch[0].to(device)\n        Y = batch[1].to(device)\n        mask = batch[2].to(device)\n        preds = model(X) * mask\n        loss = criterion(preds, Y)\n\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:48:21.982288Z","iopub.execute_input":"2024-12-01T16:48:21.98279Z","iopub.status.idle":"2024-12-01T16:48:21.989915Z","shell.execute_reply.started":"2024-12-01T16:48:21.982749Z","shell.execute_reply":"2024-12-01T16:48:21.988514Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def evaluate(model, data_loader, criterion, device):\n    model.eval()\n    n = 0\n    total_loss = 0.0\n    for batch in data_loader:\n        X = batch[0].to(device)\n        Y = batch[1].to(device)\n        mask = batch[2].to(device)\n        \n        with torch.no_grad():\n            preds = model(X) * mask\n        \n        loss = criterion(preds, Y)\n        total_loss += loss.item() * X.shape[0]\n        n += X.shape[0]\n    \n    avg_loss = total_loss / n\n    \n    return avg_loss","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:48:31.953602Z","iopub.execute_input":"2024-12-01T16:48:31.954116Z","iopub.status.idle":"2024-12-01T16:48:31.962118Z","shell.execute_reply.started":"2024-12-01T16:48:31.954059Z","shell.execute_reply":"2024-12-01T16:48:31.960397Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def predict(model, data_loader, device):\n    model.eval()\n    preds_all = []\n    for batch in data_loader:\n        X = batch[0].to(device)\n        mask = batch[2].to(device)\n        \n        with torch.no_grad():\n            preds = model(X) * mask\n            \n        preds = preds.cpu().numpy()\n        preds_all.append(preds)\n    \n    preds_all = np.concatenate(preds_all)\n        \n    return preds_all","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:48:40.086103Z","iopub.execute_input":"2024-12-01T16:48:40.087424Z","iopub.status.idle":"2024-12-01T16:48:40.095486Z","shell.execute_reply.started":"2024-12-01T16:48:40.087359Z","shell.execute_reply":"2024-12-01T16:48:40.094031Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"n_splits = 10\nn_epochs = 15\nbatch_size = 16\n    \nY = df_y.values\nmask = df_mask.values\ngroups = df_y.index.get_level_values(\"series_id\").tolist()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:48:50.245082Z","iopub.execute_input":"2024-12-01T16:48:50.24554Z","iopub.status.idle":"2024-12-01T16:48:50.252753Z","shell.execute_reply.started":"2024-12-01T16:48:50.245497Z","shell.execute_reply":"2024-12-01T16:48:50.251248Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"preds_valid = np.zeros_like(Y)\ngkf = GroupKFold(n_splits=n_splits)\nfor fold, (idx_train, idx_valid) in enumerate(gkf.split(X, Y, groups=groups)):   \n    print(f\"fold: {fold}\")\n    X_train = X[idx_train]\n    Y_train = Y[idx_train]\n    mask_train = mask[idx_train]\n\n    X_valid = X[idx_valid]\n    Y_valid = Y[idx_valid]\n    mask_valid = mask[idx_valid]\n\n    # dataset\n    ds_train = MyDataset(X_train, Y_train, mask_train)\n    ds_valid = MyDataset(X_valid, Y_valid, mask_valid)\n\n    # dataloader\n    dl_train = DataLoader(\n        ds_train, \n        batch_size=batch_size,\n        shuffle=True, \n        num_workers=0, \n        pin_memory=True,\n        drop_last=True\n    )\n    dl_valid = DataLoader(\n        ds_valid,\n        batch_size=batch_size,\n        shuffle=False, \n        num_workers=0, \n        pin_memory=True, \n        drop_last=False\n    )\n\n    # build model\n    torch.manual_seed(0)\n    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n    model = UNet1d(\n        input_channels=X.shape[1],\n        initial_channels=72,\n        initial_kernel_size=15,\n        down_channels=(72, 72, 72),\n        down_kernel_size=(12, 15, 15),\n        down_stride=(12, 9, 5),  # first element must be 12\n        res_depth=3,\n        res_kernel_size=15,\n        se_ratio=4,\n        out_kernel_size=21,\n    )\n    \n    if fold == 0:\n        print(summary(\n            model=model,\n            input_size=(batch_size, X.shape[1], X.shape[2]),\n            col_names=[\"input_size\", \"output_size\", \"num_params\"],\n            col_width=20\n        ))\n        \n    model.to(device)\n    optimizer = optim.Adam(model.parameters(), lr=1e-3)\n    scheduler = CosineAnnealingLR(optimizer, T_max=n_epochs)\n    criterion = nn.MSELoss()\n\n    best_loss = np.inf\n    for epoch in range(n_epochs):\n        train(model, dl_train, optimizer, criterion, device)\n        loss = evaluate(model, dl_valid, criterion, device)\n        scheduler.step()\n        if loss < best_loss:\n            best_loss = loss\n            torch.save(model.state_dict(), f\"model_{fold}.pth\")\n            print(f\"epoch: {epoch}\\tValid-Loss: {loss}\\tBEST\")\n        else:\n            print(f\"epoch: {epoch}\\tValid-Loss: {loss}\")\n\n    with torch.no_grad():\n        model.load_state_dict(torch.load(f\"model_{fold}.pth\"))\n        preds_valid[idx_valid] = predict(model, dl_valid, device)\n        \n    print(f\"Best_Loss: {best_loss}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-01T16:49:16.48044Z","iopub.execute_input":"2024-12-01T16:49:16.480862Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}