{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":53666,"databundleVersionId":6589269,"sourceType":"competition"},{"sourceId":6684565,"sourceType":"datasetVersion","datasetId":3714177},{"sourceId":178858118,"sourceType":"kernelVersion"},{"sourceId":178873842,"sourceType":"kernelVersion"}],"dockerImageVersionId":30559,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"Finding Critical Points using RNNs\n\n - [Data Prep notebook](https://www.kaggle.com/code/werus23/sleep-critical-point-prepare-data)\n - Training notebook - this\n - [Inference Notebook](https://www.kaggle.com/code/werus23/child-sleep-critical-point-regression)\n\nCredits:\n\n - idea: https://www.kaggle.com/competitions/child-mind-institute-detect-sleep-states/discussion/441470\n - dataloader: https://www.kaggle.com/code/henriupton/efficient-loading-memory-usage-visualizations-cmi\n - arch: https://www.kaggle.com/competitions/tlvmc-parkinsons-freezing-gait-prediction/discussion/416410","metadata":{}},{"cell_type":"markdown","source":"# 1. Configuration","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport gc\nimport time\nimport json\nfrom datetime import datetime\nimport matplotlib.pyplot as plt\nimport os\nimport joblib\nimport random\nimport math\nfrom tqdm.auto import tqdm \nfrom scipy.stats import mode \nfrom scipy.interpolate import interp1d\n\nfrom math import pi, sqrt, exp\nimport sklearn,sklearn.model_selection\nimport torch\nfrom torch import nn,Tensor\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader, Dataset, SubsetRandomSampler\nfrom sklearn.metrics import average_precision_score\nfrom timm.scheduler import CosineLRScheduler\nplt.style.use(\"ggplot\")\n\nfrom pyarrow.parquet import ParquetFile\nimport pyarrow as pa \nimport ctypes\ntorch.manual_seed(1234)\ndef normalize(y):\n    mean = y[:,0].mean().item()\n    std = y[:,0].std().item()\n    y[:,0] = (y[:,0]-mean)/(std+1e-16)\n    mean = y[:,1].mean().item()\n    std = y[:,1].std().item()\n    y[:,1] = (y[:,1]-mean)/(std+1e-16)\n    return y\n# def normalize(y):\n#     min_val = y[:,0].min().item()  # Convert to NumPy array and extract the value\n#     max_val = y[:,0].max().item()  # Convert to NumPy array and extract the value\n#     y[:,0]= (y[:,0] - min_val) / (max_val - min_val)\n#     min_val = y[:,1].min().item()  # Convert to NumPy array and extract the value\n#     max_val = y[:,1].max().item()  # Convert to NumPy array and extract the value\n#     y[:,1]= (y[:,1] - min_val) / (max_val - min_val)\n#     return y\n\n\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-21T14:34:29.341025Z","iopub.execute_input":"2024-05-21T14:34:29.341393Z","iopub.status.idle":"2024-05-21T14:34:29.352755Z","shell.execute_reply.started":"2024-05-21T14:34:29.341364Z","shell.execute_reply":"2024-05-21T14:34:29.351652Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"EPOCHS =5\nWARMUP_PROP = 0.2\nBS = 1\nWORKERS = 4\nTRAIN_PROP = 0.9\nmax_chunk_size = 24*60*100\nif device=='cpu':\n    torch.set_num_interop_threads(WORKERS)\n    torch.set_num_threads(WORKERS)","metadata":{"execution":{"iopub.status.busy":"2024-05-21T14:06:02.409635Z","iopub.execute_input":"2024-05-21T14:06:02.410032Z","iopub.status.idle":"2024-05-21T14:06:02.415620Z","shell.execute_reply.started":"2024-05-21T14:06:02.410000Z","shell.execute_reply":"2024-05-21T14:06:02.414527Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def plot_history(history, model_path=\".\", show=True):\n    epochs = range(1, len(history[\"train_loss\"]) + 1)\n\n    plt.figure()\n    plt.plot(epochs, history[\"train_loss\"], label=\"Training Loss\")\n    plt.plot(epochs, history[\"valid_loss\"], label=\"Validation Loss\")\n    plt.title(\"Loss evolution\")\n    plt.xlabel(\"Epochs\")\n    plt.ylabel(\"Loss\")\n    plt.legend()\n    plt.savefig(os.path.join(model_path, \"loss_evo.png\"))\n    if show:\n        plt.show()\n    plt.close()\n\n#     plt.figure()\n#     plt.plot(epochs, history[\"valid_mAP\"])\n#     plt.title(\"Validation mAP evolution\")\n#     plt.xlabel(\"Epochs\")\n#     plt.ylabel(\"mAP\")\n#     plt.savefig(os.path.join(model_path, \"mAP_evo.png\"))\n#     if show:\n#         plt.show()\n#     plt.close()\n\n    plt.figure()\n    plt.plot(epochs, history[\"lr\"])\n    plt.title(\"Learning Rate evolution\")\n    plt.xlabel(\"Epochs\")\n    plt.ylabel(\"LR\")\n    plt.savefig(os.path.join(model_path, \"lr_evo.png\"))\n    if show:\n        plt.show()\n    plt.close()","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-21T14:04:08.367376Z","iopub.execute_input":"2024-05-21T14:04:08.367688Z","iopub.status.idle":"2024-05-21T14:04:08.375952Z","shell.execute_reply.started":"2024-05-21T14:04:08.367662Z","shell.execute_reply":"2024-05-21T14:04:08.375103Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ResidualBiGRU(nn.Module):\n    def __init__(self, hidden_size, n_layers=1, bidir=True):\n        super(ResidualBiGRU, self).__init__()\n\n        self.hidden_size = hidden_size\n        self.n_layers = n_layers\n\n        self.gru = nn.GRU(\n            hidden_size,\n            hidden_size,\n            n_layers,\n            batch_first=True,\n            bidirectional=bidir,\n        )\n        dir_factor = 2 if bidir else 1\n        self.fc1 = nn.Linear(\n            hidden_size * dir_factor, hidden_size * dir_factor * 2\n        )\n        self.ln1 = nn.LayerNorm(hidden_size * dir_factor * 2)\n        self.fc2 = nn.Linear(hidden_size * dir_factor * 2, hidden_size)\n        self.ln2 = nn.LayerNorm(hidden_size)\n\n    def forward(self, x, h=None):\n        res, new_h = self.gru(x, h)\n        # res.shape = (batch_size, sequence_size, 2*hidden_size)\n\n        res = self.fc1(res)\n        res = self.ln1(res)\n        res = nn.functional.relu(res)\n\n        res = self.fc2(res)\n        res = self.ln2(res)\n        res = nn.functional.relu(res)\n\n        # skip connection\n        res = res + x\n\n        return res, new_h\n\nclass MultiResidualBiGRU(nn.Module):\n    def __init__(self, input_size, hidden_size, out_size, n_layers, bidir=True):\n        super(MultiResidualBiGRU, self).__init__()\n\n        self.input_size = input_size\n        self.hidden_size = hidden_size\n        self.out_size = out_size\n        self.n_layers = n_layers\n\n        self.fc_in = nn.Linear(input_size, hidden_size)\n        self.ln = nn.LayerNorm(hidden_size)\n        self.res_bigrus = nn.ModuleList(\n            [\n                ResidualBiGRU(hidden_size, n_layers=1, bidir=bidir)\n                for _ in range(n_layers)\n            ]\n        )\n        self.fc_out = nn.Linear(hidden_size, out_size)\n\n    def forward(self, x, h=None):\n        # if we are at the beginning of a sequence (no hidden state)\n        if h is None:\n            # (re)initialize the hidden state\n            h = [None for _ in range(self.n_layers)]\n\n        x = self.fc_in(x)\n        x = self.ln(x)\n        x = nn.functional.relu(x)\n\n        new_h = []\n        for i, res_bigru in enumerate(self.res_bigrus):\n            x, new_hi = res_bigru(x, h[i])\n            new_h.append(new_hi)\n\n        x = self.fc_out(x)\n#         x = F.normalize(x,dim=0)\n        return x, new_h  # log probabilities + hidden states","metadata":{"_kg_hide-output":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-21T14:04:08.377001Z","iopub.execute_input":"2024-05-21T14:04:08.377273Z","iopub.status.idle":"2024-05-21T14:04:08.396642Z","shell.execute_reply.started":"2024-05-21T14:04:08.377251Z","shell.execute_reply":"2024-05-21T14:04:08.395900Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Define Dataset","metadata":{}},{"cell_type":"code","source":"SIGMA = 720 #average length of day is 24*60*12 = 17280 for comparison\nSAMPLE_FREQ = 12 # 1 obs per minute\nfrom sklearn.preprocessing import MinMaxScaler\nclass SleepDataset(Dataset):\n    def __init__(self,file):\n        self.targets,self.data,self.ids = joblib.load(file)\n        self.scaler = MinMaxScaler()\n            \n    def downsample_seq_generate_features(self,feat, downsample_factor = SAMPLE_FREQ):\n        # downsample data and generate features\n        if len(feat)%SAMPLE_FREQ==0:\n            feat = np.concatenate([feat,np.zeros(SAMPLE_FREQ-((len(feat))%SAMPLE_FREQ))+feat[-1]])\n        feat = np.reshape(feat, (-1,SAMPLE_FREQ))\n        feat_mean = np.mean(feat,1)\n        feat_std = np.std(feat,1)\n        feat_median = np.median(feat,1)\n        feat_max = np.max(feat,1)\n        feat_min = np.min(feat,1)\n\n        return np.dstack([feat_mean,feat_std,feat_median,feat_max,feat_min])[0]\n    def downsample_seq(self,feat, downsample_factor = SAMPLE_FREQ):\n        # downsample data\n        if len(feat)%SAMPLE_FREQ==0:\n            feat = np.concatenate([feat,np.zeros(SAMPLE_FREQ-((len(feat))%SAMPLE_FREQ))+feat[-1]])\n        feat = np.reshape(feat, (-1,SAMPLE_FREQ))\n        feat_mean = np.mean(feat,1)\n        return feat_mean\n    \n    def downsample_categorical_mode(self, cat_feat, downsample_factor=SAMPLE_FREQ):\n        if len(cat_feat) % SAMPLE_FREQ == 0:\n            cat_feat = np.concatenate([cat_feat,np.zeros(SAMPLE_FREQ-((len(cat_feat))%SAMPLE_FREQ))+cat_feat[-1]])\n        cat_feat = np.reshape(cat_feat, (-1, SAMPLE_FREQ))\n        cat_mode = mode(cat_feat, axis=1)[0]\n        return cat_mode\n    \n    def gauss(self,n=SIGMA,sigma=SIGMA*0.15):\n        # guassian distribution function\n        r = range(-int(n/2),int(n/2)+1)\n        return [1 / (sigma * sqrt(2*pi)) * exp(-float(x)**2/(2*sigma**2)) for x in r]\n    \n    def __len__(self):\n        return len(self.targets)\n\n    def __getitem__(self, index):\n        X = self.data[index][['anglez','enmo']]\n        y = self.targets[index]\n#         cat_feats = self.data[index][['hour', 'minute', 'weekday']]\n        # turn target inds into array\n#         cat_feats_downsampled = np.concatenate([self.downsample_categorical_mode(cat_feats.values[:, i], SAMPLE_FREQ)for i in range(cat_feats.shape[1])],-1)\n#         print(cat_feats_downsampled)\n        target_guassian = np.zeros((len(X),2))\n        for s,e in y:\n            st1,st2 = max(0,s-SIGMA//2),s+SIGMA//2+1\n            ed1,ed2 = e-SIGMA//2,min(len(X),e+SIGMA//2+1)\n            target_guassian[st1:st2,0] = self.gauss()[st1-(s-SIGMA//2):]\n            target_guassian[ed1:ed2,1] = self.gauss()[:SIGMA+1-((e+SIGMA//2+1)-ed2)]\n            gc.collect()\n        y = target_guassian\n        gc.collect()\n        X = np.stack([self.downsample_seq(X.values[:,i],SAMPLE_FREQ) for i in range(X.shape[1])],1)\n        X = self.scaler.fit_transform(X)\n        gc.collect()\n#         print(X.shape,downsample_seq.shape,cat_feats_downsampled.shape)\n#         X = np.concatenate([downsample_seq, cat_feats_downsampled], axis=1)\n        gc.collect()\n        y = np.dstack([self.downsample_seq(y[:,i],SAMPLE_FREQ) for i in range(y.shape[1])])[0]\n        gc.collect()\n        y = normalize(torch.from_numpy(y))\n        X = torch.from_numpy(X)\n        return X, y\ntrain_ds = SleepDataset('/kaggle/input/k/tlhmeomeo/sleep-critical-point-prepare-data/train_data.pkl')","metadata":{"execution":{"iopub.status.busy":"2024-05-21T14:04:08.398847Z","iopub.execute_input":"2024-05-21T14:04:08.399157Z","iopub.status.idle":"2024-05-21T14:04:35.815902Z","shell.execute_reply.started":"2024-05-21T14:04:08.399133Z","shell.execute_reply":"2024-05-21T14:04:35.814890Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"SIGMA = 720 #average length of day is 24*60*12 = 17280 for comparison\nSAMPLE_FREQ = 12 # 1 obs per minute\nfrom sklearn.preprocessing import MinMaxScaler\nclass SleepDataset(Dataset):\n    def __init__(self,file):\n        self.targets,self.data,self.ids = joblib.load(file)\n        self.scaler = MinMaxScaler()\n            \n            \n    def downsample_seq_generate_features(self,feat, downsample_factor = SAMPLE_FREQ):\n        # downsample data and generate features\n        if len(feat)%SAMPLE_FREQ==0:\n            feat = np.concatenate([feat,np.zeros(SAMPLE_FREQ-((len(feat))%SAMPLE_FREQ))+feat[-1]])\n        feat = np.reshape(feat, (-1,SAMPLE_FREQ))\n        feat_mean = np.mean(feat,1)\n        feat_std = np.std(feat,1)\n        feat_median = np.median(feat,1)\n        feat_max = np.max(feat,1)\n        feat_min = np.min(feat,1)\n\n        return np.dstack([feat_mean,feat_std,feat_median,feat_max,feat_min])[0]\n    def downsample_seq(self,feat, downsample_factor = SAMPLE_FREQ):\n        # downsample data\n        if len(feat)%SAMPLE_FREQ==0:\n            feat = np.concatenate([feat,np.zeros(SAMPLE_FREQ-((len(feat))%SAMPLE_FREQ))+feat[-1]])\n        feat = np.reshape(feat, (-1,SAMPLE_FREQ))\n        feat_mean = np.mean(feat,1)\n        return feat_mean\n    \n    def downsample_categorical_mode(self, cat_feat, downsample_factor=SAMPLE_FREQ):\n        if len(cat_feat) % SAMPLE_FREQ == 0:\n            cat_feat = np.concatenate([cat_feat,np.zeros(SAMPLE_FREQ-((len(cat_feat))%SAMPLE_FREQ))+cat_feat[-1]])\n        cat_feat = np.reshape(cat_feat, (-1, SAMPLE_FREQ))\n        cat_mode = mode(cat_feat, axis=1)[0]\n        return cat_mode\n    \n    def gauss(self,n=SIGMA,sigma=SIGMA*0.15):\n        # guassian distribution function\n        r = range(-int(n/2),int(n/2)+1)\n        return [1 / (sigma * sqrt(2*pi)) * exp(-float(x)**2/(2*sigma**2)) for x in r]\n    \n    def __len__(self):\n        return len(self.targets)\n\n    def __getitem__(self, index):\n        X = self.data[index][['anglez','enmo']]\n        y = self.targets[index]\n        cat_feats = self.data[index][['hour', 'minute', 'weekday']]\n        # turn target inds into array\n        cat_feats_downsampled = np.concatenate([self.downsample_categorical_mode(cat_feats.values[:, i], SAMPLE_FREQ)for i in range(cat_feats.shape[1])],-1)\n#         print(cat_feats_downsampled)\n        target_guassian = np.zeros((len(X),2))\n        for s,e in y:\n            st1,st2 = max(0,s-SIGMA//2),s+SIGMA//2+1\n            ed1,ed2 = e-SIGMA//2,min(len(X),e+SIGMA//2+1)\n            target_guassian[st1:st2,0] = self.gauss()[st1-(s-SIGMA//2):]\n            target_guassian[ed1:ed2,1] = self.gauss()[:SIGMA+1-((e+SIGMA//2+1)-ed2)]\n            gc.collect()\n        y = target_guassian\n        gc.collect()\n        downsample_seq = np.concatenate([self.downsample_seq_generate_features(X.values[:,i],SAMPLE_FREQ) for i in range(X.shape[1])],-1)\n        gc.collect()\n#         print(X.shape,downsample_seq.shape,cat_feats_downsampled.shape)\n\n        X = np.concatenate([downsample_seq, cat_feats_downsampled], axis=1)\n#         X = self.scaler.fit_transform(X)\n        gc.collect()\n        y = np.dstack([self.downsample_seq(y[:,i],SAMPLE_FREQ) for i in range(y.shape[1])])[0]\n        gc.collect()\n        y = normalize(torch.from_numpy(y))\n        X = torch.from_numpy(X)\n        return X, y\ntrain_ds = SleepDataset('/kaggle/input/k/tlhmeomeo/sleep-critical-point-prepare-data/train_data.pkl')","metadata":{"execution":{"iopub.status.busy":"2024-05-21T14:34:36.391463Z","iopub.execute_input":"2024-05-21T14:34:36.392326Z","iopub.status.idle":"2024-05-21T14:34:40.122807Z","shell.execute_reply.started":"2024-05-21T14:34:36.392290Z","shell.execute_reply":"2024-05-21T14:34:40.122035Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train and Eval","metadata":{}},{"cell_type":"code","source":"train_size = int(TRAIN_PROP * len(train_ds))\nvalid_size = len(train_ds) - train_size\nindices = torch.randperm(len(train_ds))\ntrain_sampler = SubsetRandomSampler(indices[:train_size])\nvalid_sampler = SubsetRandomSampler(\n    indices[train_size : train_size + valid_size]\n)\nsteps = train_size*EPOCHS\nwarmup_steps = int(steps*WARMUP_PROP)\nmodel = MultiResidualBiGRU(input_size=13,hidden_size=64,out_size=2,n_layers=5).to(device) \nmodel.load_state_dict(torch.load('/kaggle/input/sleep-critical-point-train/model_best.pth'))\n\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3,weight_decay = 0)\nscheduler = CosineLRScheduler(optimizer,t_initial= steps,warmup_t=warmup_steps, warmup_lr_init=1e-6,lr_min=2e-8,)\ndt = time.time()\nmodel_path = '/kaggle/working/'\nhistory = {\n    \"train_loss\": [],\n    \"valid_loss\": [],\n    \"valid_mAP\": [],\n    \"lr\": [],\n}\n\nbest_valid_loss = np.inf\ncriterion = torch.nn.MSELoss()","metadata":{"execution":{"iopub.status.busy":"2024-05-21T14:06:06.440309Z","iopub.execute_input":"2024-05-21T14:06:06.441028Z","iopub.status.idle":"2024-05-21T14:06:06.840414Z","shell.execute_reply.started":"2024-05-21T14:06:06.440994Z","shell.execute_reply":"2024-05-21T14:06:06.839612Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for i in range(1):\n    sample = train_ds[i]\n    print(\"Data\",sample[0])\n    print(\"target\",sample[1])","metadata":{"execution":{"iopub.status.busy":"2024-05-21T14:06:12.550647Z","iopub.execute_input":"2024-05-21T14:06:12.551012Z","iopub.status.idle":"2024-05-21T14:06:19.153690Z","shell.execute_reply.started":"2024-05-21T14:06:12.550981Z","shell.execute_reply":"2024-05-21T14:06:19.152782Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def evaluate(\n    model: nn.Module, max_chunk_size: int, loader: DataLoader, device, criterion\n):\n    model.eval()\n    valid_loss = 0.0\n    y_true_full = torch.FloatTensor([]).half()\n    y_pred_full = torch.FloatTensor([]).half()\n    for X_batch, y_batch in tqdm(loader, desc=\"Eval\", unit=\"batch\"):\n        # as some of the sequences we are dealing with are pretty long, we\n        # use a chunk-based approach\n        y_batch = y_batch.to(device, non_blocking=True)\n        pred = torch.zeros(y_batch.shape).to(device, non_blocking=True).half()\n\n        # (re)initialize model's hidden state(s)\n        h = None\n\n        # number of chunks for this sequence (we assume batch size = 1)\n        seq_len = X_batch.shape[1]\n        for i in range(0, seq_len, max_chunk_size):\n            X_chunk = X_batch[:, i : i + max_chunk_size].float().to(device, non_blocking=True)\n\n            y_pred, h = model(X_chunk, h)\n            h = [hi.detach() for hi in h]\n            pred[:, i : i + max_chunk_size] = y_pred.half()\n            del X_chunk\n            gc.collect()\n        loss = criterion(\n            pred.float(),\n            y_batch.float(),\n        )\n        valid_loss += loss.item()\n        del pred,loss\n        gc.collect()\n\n    valid_loss /= len(loader)\n\n    y_true_full = y_true_full.squeeze(0)\n    y_pred_full = y_pred_full.squeeze(0)\n    gc.collect()\n    return valid_loss","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2024-05-21T13:22:25.746685Z","iopub.status.idle":"2024-05-21T13:22:25.747105Z","shell.execute_reply.started":"2024-05-21T13:22:25.746896Z","shell.execute_reply":"2024-05-21T13:22:25.746916Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_loader = DataLoader(\n    train_ds,\n    batch_size=BS,\n    sampler=train_sampler,\n    pin_memory=True,\n    num_workers=1,\n)\nvalid_loader = DataLoader(\n    train_ds,\n    batch_size=1,\n    sampler=valid_sampler,\n    pin_memory=True,\n    num_workers=1,\n)\n# for epoch in range(1, EPOCHS + 1):\n#     train_loss = 0.0\n#     n_tot_chunks = 0\n#     pbar = tqdm(\n#         train_loader, desc=\"Training\", unit=\"batch\"\n#     )\n#     model.train()\n#     for step,(X_batch, y_batch) in enumerate(pbar):\n#         y_batch = y_batch.to(device, non_blocking=True)\n#         pred = torch.zeros(y_batch.shape).to(device, non_blocking=True)\n#         optimizer.zero_grad()\n#         scheduler.step(step+train_size*epoch)\n#         h = None\n\n#         seq_len = X_batch.shape[1]\n#         for i in range(0, seq_len, max_chunk_size):\n#             X_chunk = X_batch[:, i : i + max_chunk_size].float()\n\n#             X_chunk = X_chunk.to(device, non_blocking=True)\n\n#             y_pred, h = model(X_chunk, h)\n#             h = [hi.detach() for hi in h]\n#             pred[:, i : i + max_chunk_size] = y_pred\n#             del X_chunk,y_pred\n\n#         loss = criterion(\n#             normalize(pred).float(),\n#             y_batch.float(),\n#         )\n#         loss.backward()\n#         train_loss += loss.item()\n#         n_tot_chunks+=1\n#         pbar.set_description(f'Training: loss = {(train_loss/n_tot_chunks):.2f}')\n\n#         nn.utils.clip_grad_norm_(model.parameters(), max_norm=1e-1)\n#         optimizer.step()\n#         del pred,loss,y_batch,X_batch,h\n#         gc.collect()\n#     train_loss /= len(train_loader)\n#     del pbar\n#     gc.collect()\n#     ctypes.CDLL(\"libc.so.6\").malloc_trim(0)\n\n#     if epoch % 1 == 0:\n#         valid_loss = evaluate(\n#             model, max_chunk_size, valid_loader, device, criterion\n#         )\n\n#         history[\"train_loss\"].append(train_loss)\n#         history[\"valid_loss\"].append(valid_loss)\n#         history[\"lr\"].append(optimizer.param_groups[0][\"lr\"])\n\n#         if valid_loss < best_valid_loss:\n#             best_valid_loss = valid_loss\n#             torch.save(\n#                 model.state_dict(),\n#                 os.path.join(model_path, f\"model_best.pth\"),\n#             )\n\n#         dt = time.time() - dt\n#         print(\n#             f\"{epoch}/{EPOCHS} -- \",\n#             f\"train_loss = {train_loss:.6f} -- \",\n#             f\"valid_loss = {valid_loss:.6f} -- \",\n#             f\"time = {dt:.6f}s\",\n#         )\n#         dt = time.time()\n\n# plot_history(history, model_path=model_path)\n# history_path = os.path.join(model_path, \"history.json\")\n# with open(history_path, \"w\", encoding=\"utf-8\") as f:\n#     json.dump(history, f, ensure_ascii=False, indent=4)","metadata":{"execution":{"iopub.status.busy":"2024-05-21T14:04:40.214307Z","iopub.status.idle":"2024-05-21T14:04:40.214615Z","shell.execute_reply.started":"2024-05-21T14:04:40.214463Z","shell.execute_reply":"2024-05-21T14:04:40.214477Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\ndef visualize_sample(model, dataset, device, num_samples=2, max_chunk_size=300):\n    model.eval()\n    \n    # Keep track of the number of samples visualized\n    samples_visualized = 0\n    \n    # Iterate through the dataset to obtain samples\n    for sample_index in range(len(dataset)):\n        # Get the sample\n        X_sample, y_sample = dataset[sample_index]\n        \n        # Move data to the device\n        X_sample = X_sample.unsqueeze(0).to(device)  # Add batch dimension\n        y_sample = y_sample.unsqueeze(0).to(device)  # Add batch dimension\n        \n        # Initialize model's hidden state\n        h = None\n        \n        # Number of chunks for this sequence (assuming batch size = 1)\n        seq_len = X_sample.shape[1]\n        \n        # Iterate through the sequence chunks\n        for i in range(0, seq_len, max_chunk_size):\n            X_chunk = X_sample[:, i : i + max_chunk_size].float().to(device, non_blocking=True)\n            y_chunk = y_sample[:, i : i + max_chunk_size]\n            y_pred, h = model(X_chunk, h)\n            y_pred = y_pred.half().float()  # Convert to float\n\n            h = [hi.detach() for hi in h]\n            \n            # Plot the features, true labels, and predictions for the current chunk\n            plot_features_and_labels(X_chunk, y_chunk, y_pred)\n        \n        # Increment the count of samples visualized\n        samples_visualized += 1\n        \n        # Break the loop if the desired number of samples is reached\n        if samples_visualized == num_samples:\n            break\n\ndef plot_features_and_labels(X_chunk, y_chunk, y_pred):\n    fig, ax1 = plt.subplots(figsize=(20, 8))\n    \n    # Plot feature 4\n    ax1.plot(X_chunk[0, :, 4].cpu().numpy(), color='tab:blue', label='Feature 4')\n    ax1.set_ylabel('anglez', color='tab:blue')\n    ax1.tick_params(axis='y', labelcolor='tab:blue')\n    \n    # Create another y-axis for feature 9\n    ax2 = ax1.twinx()\n    ax2.plot(X_chunk[0, :, 9].cpu().numpy(), color='tab:red', label='Feature 9')\n    ax2.set_ylabel('enmo', color='tab:red')\n    ax2.tick_params(axis='y', labelcolor='tab:red')\n    \n    # Plot true labels and predictions for y[0]\n    ax1.plot(y_chunk[0, :, 0].cpu().detach().numpy(), label='True Onset', linestyle='--', color='tab:green')\n    ax1.plot(y_pred[0, :, 0].cpu().detach().numpy(), label='Prediction Onset', linestyle=':', color='tab:green')\n    ax1.plot(y_chunk[0, :, 1].cpu().detach().numpy(), label='True Label Wakeup', linestyle='--', color='orange')\n    ax1.plot(y_pred[0, :, 1].cpu().detach().numpy(), label='Prediction Wakeup', linestyle=':', color='orange')\n    ax1.set_xlabel('Time')\n    \n    # Combine legends\n    lines, labels = ax1.get_legend_handles_labels()\n    lines2, labels2 = ax2.get_legend_handles_labels()\n    ax2.legend(lines + lines2, labels + labels2, loc='upper right')\n    \n    plt.show()\n\n\n\n# Call the visualize_sample method\nvisualize_sample(model, train_ds, device, num_samples=1)\n","metadata":{"execution":{"iopub.status.busy":"2024-05-21T14:56:41.737022Z","iopub.execute_input":"2024-05-21T14:56:41.737405Z","iopub.status.idle":"2024-05-21T14:57:39.744557Z","shell.execute_reply.started":"2024-05-21T14:56:41.737374Z","shell.execute_reply":"2024-05-21T14:57:39.743560Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}