{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"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 \n\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\n\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\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":"2023-10-14T02:42:33.412162Z","iopub.execute_input":"2023-10-14T02:42:33.412881Z","iopub.status.idle":"2023-10-14T02:42:40.549168Z","shell.execute_reply.started":"2023-10-14T02:42:33.412848Z","shell.execute_reply":"2023-10-14T02:42:40.548169Z"},"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 = 150000\nif device=='cpu':\n    torch.set_num_interop_threads(WORKERS)\n    torch.set_num_threads(WORKERS)","metadata":{"execution":{"iopub.status.busy":"2023-10-14T02:42:40.554390Z","iopub.execute_input":"2023-10-14T02:42:40.556571Z","iopub.status.idle":"2023-10-14T02:42:40.563485Z","shell.execute_reply.started":"2023-10-14T02:42:40.556536Z","shell.execute_reply":"2023-10-14T02:42:40.562479Z"},"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":"2023-10-14T02:42:40.564608Z","iopub.execute_input":"2023-10-14T02:42:40.568751Z","iopub.status.idle":"2023-10-14T02:42:40.579907Z","shell.execute_reply.started":"2023-10-14T02:42:40.568720Z","shell.execute_reply":"2023-10-14T02:42:40.578912Z"},"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":"2023-10-14T02:42:40.583947Z","iopub.execute_input":"2023-10-14T02:42:40.586243Z","iopub.status.idle":"2023-10-14T02:42:40.603219Z","shell.execute_reply.started":"2023-10-14T02:42:40.586213Z","shell.execute_reply":"2023-10-14T02:42:40.602152Z"},"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\nclass SleepDataset(Dataset):\n    def __init__(\n        self,\n        file\n    ):\n        self.targets,self.data,self.ids = joblib.load(file)\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 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        \n        # turn target inds into array\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.concatenate([self.downsample_seq_generate_features(X.values[:,i],SAMPLE_FREQ) for i in range(X.shape[1])],-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/sleep-critical-point-prepare-data/train_data.pkl')","metadata":{"execution":{"iopub.status.busy":"2023-10-14T02:42:40.608745Z","iopub.execute_input":"2023-10-14T02:42:40.609035Z","iopub.status.idle":"2023-10-14T02:42:49.444521Z","shell.execute_reply.started":"2023-10-14T02:42:40.609014Z","shell.execute_reply":"2023-10-14T02:42:49.443396Z"},"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=10,hidden_size=64,out_size=2,n_layers=5).to(device)\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}\nbest_valid_loss = np.inf\ncriterion = torch.nn.MSELoss()","metadata":{"execution":{"iopub.status.busy":"2023-10-14T02:42:49.446053Z","iopub.execute_input":"2023-10-14T02:42:49.446719Z","iopub.status.idle":"2023-10-14T02:42:54.723031Z","shell.execute_reply.started":"2023-10-14T02:42:49.446684Z","shell.execute_reply":"2023-10-14T02:42:54.722128Z"},"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":"2023-10-14T02:42:54.724439Z","iopub.execute_input":"2023-10-14T02:42:54.724747Z","iopub.status.idle":"2023-10-14T02:42:54.733227Z","shell.execute_reply.started":"2023-10-14T02:42:54.724718Z","shell.execute_reply":"2023-10-14T02:42:54.732199Z"},"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)\nfor 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\nplot_history(history, model_path=model_path)\nhistory_path = os.path.join(model_path, \"history.json\")\nwith open(history_path, \"w\", encoding=\"utf-8\") as f:\n    json.dump(history, f, ensure_ascii=False, indent=4)","metadata":{"execution":{"iopub.status.busy":"2023-10-14T02:42:54.734826Z","iopub.execute_input":"2023-10-14T02:42:54.735434Z"},"trusted":true},"execution_count":null,"outputs":[]}]}