{"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":"# Pytorch-lightning + LSTM\n\nFollowing is an example how one can rewrite pytorch LSTM model in a Lightning form. The goal is to create 'trainer' instance from Trainer class of pytorch lightning. A ```Trainer``` consists of ```System``` part and ```DataModule``` part. ```System``` part consists of ```Model``` part (including forward method) and training part (including loss function, optimizer etc.). ```DataModule``` consists of ```Dataloader``` for training, validation and test dataset. The datasets are defined by pytorch ```Dataset``` class. In short, the relations between each class are \n\n\n`\nTrainer --+-- System -----+- Model (network definition and forward method)\n          |               | \n          |               +- training scheme (optimizer, loss function etc.)\n          | \n          +-- Data Module -- Dataloader\n                             +-- DataSet       \n`\n\n\nIn the example below, the components are named as follows. \n\n```\nIndoorDataset3    : Dataset\nIndoorLSTM       : Model \nIndoorDataModule : DataModule\nIndoorSystem     : System to feed to trainer (includes training system inside)\n```\n\n## Sumamry\n\n- Lightning system can be defined by Model + training scheme separately (in most examples available in the internet, ```Model``` (=network) definition and the training scheme are put in one big ```System``` class). \n\n- I tried to perform inference, either using ```trainer.test()``` or ```trainer.predict()```. Somehow neither of them works, and I ended up using a callback to store the best model during the training, and reload it specifically, for the inference. There must be a smarter way. \n\n\n## Acknowledgement\n\nI used the pytorch LSTM model created by @captainqxy. <br>\n**LSTM by pytorch with Unified Wi-Fi Feats** <br>\nhttps://www.kaggle.com/luffy521/lstm-by-pytorch-with-unified-wi-fi-feats\n\nThe model above is based on <br>\n**LSTM by Keras with Unified Wi-Fi Feats** by @kokitanisaka <br>\nhttps://www.kaggle.com/kokitanisaka/lstm-by-keras-with-unified-wi-fi-feats\n\nI used the data published by @kokitanisaka <br>\n**Make dataset with Wi-Fi and Beacon** <br>\nhttps://www.kaggle.com/kokitanisaka/make-dataset-with-wi-fi-and-beacon\n\nThank you so much, @captainqxy and  @kokitanisaka.","metadata":{}},{"cell_type":"markdown","source":"Here the packages are imported.","metadata":{}},{"cell_type":"code","source":"import os\nfrom pathlib import Path\nimport warnings\nfrom pytorch_lightning import Trainer\nfrom pytorch_lightning.callbacks import ModelCheckpoint\nfrom pytorch_lightning.core.datamodule import LightningDataModule\nfrom pytorch_lightning.core.lightning import LightningModule\n\nfrom torch.utils.data import Dataset, DataLoader\nimport torch.nn as nn\nimport torch\nfrom sklearn.model_selection import GroupKFold\nfrom sklearn.preprocessing import StandardScaler, LabelEncoder\nfrom datatable import (dt, f, join)\nimport gc\n\nimport numpy as np\nimport pandas as pd\nimport logging\n\nimport random","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Make logging a bit quiet. Set device. Do some housekeeping stuff. ","metadata":{}},{"cell_type":"code","source":"#---------------------------------------------------------\n#logging.getLogger('pytorch_lightning').setLevel(logging.CRITICAL)\nlogging.getLogger('lightning').setLevel(logging.CRITICAL)\nwarnings.filterwarnings('ignore')\n# ========================================================\ngc.enable()\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(device)\n\npd.options.display.max_rows = 999\npd.options.display.max_colwidth = 999\ndt.options.display.max_nrows = 999\n\n# =========================================================\nSEED = 2021","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Here come the dataset and the data loader. They are combined in ```DataModule``` of Lightning. ","metadata":{}},{"cell_type":"code","source":"\n# =========================================================\n# data set / data module / data loader / model / system\n# =========================================================\n\n\nclass IndoorDataset3(Dataset):\n    def __init__(self, data, N_FEAT, flag='TRAIN'):\n        self.data = data\n        self.n_feat = N_FEAT\n        self.flag = flag\n\n    def __len__(self):\n        return self.data.shape[0]\n\n    def __getitem__(self, index):\n\n        db = self.data[index, : self.n_feat].astype(np.long)\n        dr = self.data[index, self.n_feat: (\n            self.n_feat * 2)].astype(np.float32)\n        d_xy = self.data[index, (self.n_feat * 2): -1].astype(np.float32)\n        d_floor = self.data[index, -1].astype(np.float32)\n\n        if self.flag == 'TRAIN':\n            return db, dr, d_xy, d_floor\n        else:\n            return db, dr\n\n# ============================================\n\nclass IndoorDataModule(LightningDataModule):\n\n    def __init__(self, data, test_data, tix, vix, bx, rx, N_FEAT, BS):\n        super().__init__()\n\n        self.data = data\n        self.test_data = test_data\n        self.tix = tix\n        self.vix = vix\n\n        self.bx = bx\n        self.rx = rx\n\n        self.N_FEAT = N_FEAT\n        self.BS = BS\n\n        if Path('.').cwd() == Path('/home/meg/k5'):\n            self.num_cores = multiprocessing.cpu_count()\n        else:\n            self.num_cores = 0\n\n    def prepare_data(self):\n\n        self.data_npy = self.data[:, self.bx +\n                                  self.rx+['x', 'y', 'floor']].to_numpy()\n        self.test_data_npy = self.test_data[:,\n                                            self.bx+self.rx+['x', 'y', 'floor']].to_numpy()\n\n    def train_dataloader(self):\n\n        train_ds = IndoorDataset3(\n            self.data_npy[self.tix, :], self.N_FEAT, 'TRAIN')\n        train_dl = DataLoader(train_ds, batch_size=self.BS,\n                              shuffle=True, drop_last=False, num_workers=self.num_cores)\n\n        return train_dl\n\n    def val_dataloader(self):\n        valid_ds = IndoorDataset3(\n            self.data_npy[self.vix, :], self.N_FEAT, 'TRAIN')\n        valid_dl = DataLoader(valid_ds, batch_size=self.BS,\n                              shuffle=False, drop_last=False, num_workers=self.num_cores)\n        return valid_dl\n\n    def test_dataloader(self):\n        test_ds = IndoorDataset3(self.test_data_npy, self.N_FEAT, 'TEST')\n        test_dl = DataLoader(test_ds, batch_size=self.BS,\n                             shuffle=False, num_workers=self.num_cores)\n        return test_dl\n\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"A LSTM model was copied from **LSTM by pytorch with Unified Wi-Fi Feats** <br>\nhttps://www.kaggle.com/luffy521/lstm-by-pytorch-with-unified-wi-fi-feats\n\nThe model is edited from the original to the one with 4 times smaller layers.","metadata":{}},{"cell_type":"code","source":"# ============================================\n\nclass IndoorLSTM(LightningModule):\n\n    def __init__(self, embedding_dim=64, wifi_bssids_size=16, N_FEAT=4):\n\n        super().__init__()\n\n        self.N_FEAT = N_FEAT\n\n        self.emb_BSS = nn.Embedding(wifi_bssids_size, embedding_dim)\n\n        self.lstm1 = nn.LSTM(input_size=256, hidden_size=128,\n                             dropout=0.3, bidirectional=False)\n\n        self.lstm2 = nn.LSTM(input_size=128, hidden_size=16,\n                             dropout=0.1, bidirectional=False)\n        self.lr = nn.Linear(self.N_FEAT, self.N_FEAT * embedding_dim)\n        self.lr1 = nn.Linear(self.N_FEAT * embedding_dim * 2, 256)\n        self.lr_xy = nn.Linear(16, 2)\n\n        self.lr_floor = nn.Linear(16, 1)\n        self.batch_norm1 = nn.BatchNorm1d(self.N_FEAT)\n        self.batch_norm2 = nn.BatchNorm1d(self.N_FEAT * embedding_dim * 2)\n\n        self.batch_norm3 = nn.BatchNorm1d(1)\n        self.dropout = nn.Dropout(0.3)\n\n\n    def forward(self, xb, xr):\n\n        x_bssid = self.emb_BSS(xb)\n\n        x_bssid = torch.flatten(x_bssid, start_dim=-2)\n\n        x_rssi = self.batch_norm1(xr)\n        x_rssi = self.lr(x_rssi)\n        x_rssi = torch.relu(x_rssi)\n\n        x = torch.cat([x_bssid, x_rssi], dim=-1)\n\n        x = self.batch_norm2(x)\n        x = self.dropout(x)\n\n        x = self.lr1(x)\n        x = torch.relu(x)\n\n        x = x.unsqueeze(-2)\n        x = self.batch_norm3(x)\n        x = x.transpose(0, 1)\n\n        x, _ = self.lstm1(x)\n        x = x.transpose(0, 1)\n        x = torch.relu(x)\n        x = x.transpose(0, 1)\n\n        x, _ = self.lstm2(x)\n        x = x.transpose(0, 1)\n        x = torch.relu(x)\n\n        xy = self.lr_xy(x)\n        floor = self.lr_floor(x)\n        floor = torch.relu(floor)\n\n#        return xy.squeeze(-2), floor.squeeze(-2)\n        return xy.squeeze(-2)\n\n\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The ```System``` class to feed to a Lightning trainer. The learning rate is reduced from ```LR``` to ```LR2``` by multiplying a single constant factor, ```gamma```, at each epoch. ","metadata":{}},{"cell_type":"code","source":"# ========================================================================\n\n\nclass IndoorSystem(LightningModule):\n\n    def __init__(self, model=None, fold=0, LR=0.1, LR2=1e-4, EP=16, LOG=None):\n        super().__init__()\n\n        self.model = model\n        self.fold = fold\n        self.best_score = 1000\n        self.best_loss = 1000 * 1000\n        self.best_epoch = -1\n\n        self.learning_rate = LR\n        self.LR1 = LR\n        self.LR2 = LR2\n        self.EP = EP\n\n        self.ff = LOG\n    # --------------------------------------------\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.Adam(\n            self.model.parameters(), lr=self.learning_rate, weight_decay=1e-2)\n\n        gamma = (self.LR2/self.LR1) ** (1.0 / self.EP)\n\n        scheduler = {'scheduler': torch.optim.lr_scheduler.StepLR(optimizer, 1, gamma=gamma),\n                     'interval': 'epoch'}\n\n        return [optimizer], [scheduler]\n\n    # --------------------------------------------\n    def training_step(self, batch, batch_idx):\n\n        db, dr, d_xy, _ = batch\n        xy = self.model(db, dr)\n\n        mse = nn.MSELoss()\n        loss = mse(xy, d_xy)\n\n        print(f\"\\r\\033[32mfold \\033[0m{self.fold} \", end='')\n        print(\n            f\"\\033[33mEPOCH\\033[0m{self.current_epoch: 5} \", end='')\n        print(f\"\\033[31midx\\033[0m{batch_idx: 3} \", end='')\n        print(f\"\\033[34mloss\\033[0m{loss: 10.3f} \", end='')\n\n        with open(self.ff, 'a') as ff:\n\n            print(f\"fold{self.fold} \", end='', file=ff)\n            print(f\"EPOCH{self.current_epoch:5} \", end='', file=ff)\n            print(f'idx{batch_idx:3} ', end='', file=ff)\n            print(f'loss{loss: 10.3f} ', file=ff)\n\n        return {'loss': loss}\n\n    # --------------------------------------------\n    def validation_step(self, batch, batch_idx):\n        db, dr, d_xy, d_floor = batch\n        xy = self.model(db, dr)\n\n        mse = nn.MSELoss()\n        val_loss = mse(xy, d_xy)\n\n        return {'loss': val_loss}\n\n    # --------------------------------------------\n    def training_epoch_end(self, outputs):\n        opt = self.optimizers()\n        self.learning_rate = opt.param_groups[0]['lr']\n\n    # --------------------------------------------\n    def validation_epoch_end(self, outputs):\n\n        loss = torch.stack([f['loss'] for f in outputs]).mean()\n\n        if self.best_loss > loss:\n            #            self.best_score = score\n            self.best_loss = loss\n            self.best_epoch = self.current_epoch\n\n        if self.current_epoch != 0:\n            pass\n\n#            print(f\"\\033[34mv\\033[0m{self.best_loss:10.3f} \", end='')\n#            print(f\"\\033[32mepoch\\033[0m{self.best_epoch:5} \", end='')\n#            print(f\"\\033[35mlr\\033[0m{self.learning_rate:7.4f} \", end='')\n#            print(\"\")\n\n#            with open(self.ff, 'a') as ff:\n#                print(\n#                    f'\\033[34mv\\033[0m{self.best_loss:10.3f} ', end='', file=ff)\n#                print(\n#                    f\"\\033[32mepoch\\033[0m{self.best_epoch:5} \", end='', file=ff)\n#                print(\n#                    f\"\\033[35mlr\\033[0m{self.learning_rate:7.4f} \", end='', file=ff)\n#                print(\"\", file=ff)\n         \n        self.log('val_loss', loss)\n\n\n# ========================================================================\n# class completed |  end of data module / data loader / model / system\n# ========================================================================\n\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Utilities. Note that the error in the floor predictions are not included in the metric.","metadata":{}},{"cell_type":"code","source":"\n# ============================================\n#  utilities functions\n# ============================================\n\ndef set_seed(seed=42):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n\n\ndef comp_metric2(xhat, yhat, fhat, x, y, f):\n#    intermediate = np.sqrt((xhat-x)**2 + (yhat-y)**2) + 15 * np.abs(fhat-f)\n\n    intermediate = np.sqrt((xhat-x)**2 + (yhat-y)**2)\n    return intermediate.sum()/xhat.shape[0]\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Convert the raw site-data in the form that one can feed to ```Dataset```. The training data consists of two parts: \n\n> columns of ```bx``` : the IDs of WiFi access points. They are converted to integers. <br>\ncolumns of ```rx``` : the strengths of the signals that the smart phone received from each WiFi access point. Scaled to a normal (=Gaussian) distribution (average=0, standard deviation = 1). <br>\n\nSeveral additional informations are returned from ```wifi_prep3``` in the same time. \n\n\n> ```wifi_bssids```: list of the names of the all WiFi access points installed in the building (=site). <br>\n```le```: a level encoder to convert string names of ```wifi_bssids``` to integers. <br>\n```ss```: a scaler for the WiFi signals to those in a normal distribution. The average value is subtracted from all measurements, and then divided by the standard deviation. <br>\n\n```data``` and ```test_data``` are for training (+validation) dataset and test (=inference) dataset, respectively.  Both are in the form of datatable. \n","metadata":{}},{"cell_type":"code","source":"# ========================================================\ndef wifi_prep3(data_site, test_data_site, N_FEAT):\n\n    bx = [i for i in data_site.names if i.startswith('wifi_bssid_')]\n    rx = [i for i in data_site.names if i.startswith('wifi_rssi_')]\n\n    dtmp = data_site[:, bx].copy()\n    dtmp.rbind(test_data_site[:, bx])\n\n    wifi_bssids = np.unique(dtmp.to_numpy())\n    wifi_bssids_size = len(wifi_bssids)\n\n    del dtmp\n    gc.collect()\n\n    timegapx = [i for i in data_site.names if i.startswith('wifi_timegap_')]\n    beaconx = [i for i in data_site.names if i.startswith('beacon_')]\n#    label_cols = ['site', 'path', 'timestamp', 'x', 'y', 'floor']\n    label_cols = ['site', 'path', 'x', 'y', 'floor']\n\n\n    # level encoder\n    le = LabelEncoder()\n    _ = le.fit(wifi_bssids)\n\n    ss = StandardScaler()\n    _ = ss.fit(data_site[:, rx])\n\n    data = data_site.copy()\n    data[:, rx] = ss.transform(data_site[:, rx])\n\n    for i in bx:\n        data[:, i] = le.transform(data_site[:, i])\n\n    test_data = test_data_site.copy()\n    test_data[:, rx] = ss.transform(test_data_site[:, rx])\n\n    for i in bx:\n        test_data[:, i] = le.transform(test_data_site[:, i])\n \n    # reshape\n    data = data[:, label_cols + bx + rx + timegapx + beaconx]\n    data = data[:, f[:].remove(f[bx[N_FEAT]:bx[-1]])]\n    data = data[:, f[:].remove(f[rx[N_FEAT]:rx[-1]])]\n    data = data[:, f[:].remove(f[timegapx[0]:timegapx[-1]])]\n    data = data[:, f[:].remove(f[beaconx[0]:beaconx[-1]])]\n\n#    test_data = test_data[:, label_test_cols + bx + rx + timegapx + beaconx]\n    test_data = test_data[:, label_cols + bx + rx + timegapx + beaconx]\n    test_data = test_data[:, f[:].remove(f[bx[N_FEAT]:bx[-1]])]\n    test_data = test_data[:, f[:].remove(f[rx[N_FEAT]:rx[-1]])]\n    test_data = test_data[:, f[:].remove(f[timegapx[0]:timegapx[-1]])]\n    test_data = test_data[:, f[:].remove(f[beaconx[0]:beaconx[-1]])]\n\n    bx = bx[: N_FEAT]\n    rx = rx[: N_FEAT]\n\n    return data, test_data, wifi_bssids, wifi_bssids_size, le, ss, bx, rx\n\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The main loop. For each site (=building), models are calculated using 5-fold-split training dataset. The hyperparameters (number of epochs and so on) are not optimized.","metadata":{}},{"cell_type":"code","source":"# ============================================\n#  directories\n\npath = Path('../input/unified-ds-wifi-and-beacon/')\nlog_path = Path('.')\nTRAIN = path\nTEST = path\nLOG = log_path/'log_v1.txt'\nMODEL = log_path/'model_v1'\nMODEL.mkdir(exist_ok=True)\n\n# ============================================\ntrain_list = list(sorted(TRAIN.glob('5*.csv')))\ntest_data_all = dt.fread(TEST/'test.csv')\n\nset_seed(SEED)\n# ========================================================\nN_SPLITS = 5\n\nconfig = dict(\n    N_FEAT=40,\n    BS=1024,\n    EP=8,\n    LR=0.1, LR2=1e-3)\n\n# ========================================================\ni_sx = [0, 1]  # try first 2 sites.\n\n# ========================================================\n#  Main Loop\n# ========================================================\nfor i_site, train_file in [(i, train_list[i]) for i in i_sx]:\n\n    site = train_file.stem.split('_')[0]\n\n    print(f\"\\033[35msite \\033[31m {i_site:2} \\033[0m{site}\")\n\n    with open(LOG, 'a') as ff:\n        print(f\"site {i_site:2} {site}\", file=ff)\n\n    # --------------------------------------\n    N_FEAT = config['N_FEAT']\n    BS = config['BS']\n    EP = config['EP']\n    LR = config['LR']\n    LR2 = config['LR2']\n    # --------------------------------------\n\n    data_site = dt.fread(train_file)\n    test_data_site = test_data_all[f.site == site, :]\n\n    data, test_data, wifi_bssids, wifi_bssids_size, le, ss, bx, rx = wifi_prep3(\n        data_site, test_data_site, N_FEAT)\n\n\n    gkf = GroupKFold(N_SPLITS)\n    for fold, (tix, vix) in enumerate(gkf.split(data, data[:, ['x', 'y', 'floor']],\n                                                groups=data[:, 'path'])):\n\n        dm = IndoorDataModule(data, test_data, tix, vix, bx, rx, N_FEAT, BS)\n\n        model = IndoorLSTM(\n            embedding_dim=8, wifi_bssids_size=wifi_bssids_size, N_FEAT=N_FEAT)\n\n        indoor_system = IndoorSystem(\n            model=model, fold=fold, LR=LR, LR2=LR2, EP=EP, LOG=LOG)\n\n    # ======================\n\n        checkpoint_callback = ModelCheckpoint(\n            #            monitor='val_score',\n            monitor='val_loss',\n            dirpath=MODEL,\n            filename='m-{epoch:02d}-{val_loss:.2f}',\n            save_top_k=3,\n            mode='min'\n        )\n\n        trainer = Trainer(\n            gpus=0,\n            max_epochs=EP,\n            logger=False,\n            callbacks=[checkpoint_callback],\n            checkpoint_callback=True,\n            progress_bar_refresh_rate=0)\n\n        trainer.fit(indoor_system, dm)\n        print('')\n\n    # ======================\n    #   inference\n\n        trained_model = IndoorSystem.load_from_checkpoint(\n            checkpoint_callback.best_model_path,\n            model=model, fold=fold, LR=LR, LR2=LR2, EP=EP)\n\n        _ = trained_model.eval()\n\n        xy_stack = []\n        for db, dr in dm.test_dataloader():\n            xy = trained_model.model(db, dr).cpu().detach().numpy()\n            xy_stack.append(xy)\n\n        pred_xy = np.vstack(xy_stack)\n\n    print(pred_xy.shape)\n\n# ========================================================\n\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}