{"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":"# Transformer理解のためのnotebook\ntransformerを初めて動かして理解したいという方向けのnotebookなので、精度は追求しません。\n\n精度の高いnotebookは, 公開されているnotebookをtransformerで検索すれば見つかるはずです。 \n\nTransformerの説明は[図で理解するTransformer](https://qiita.com/birdwatcher/items/b3e4428f63f708db37b7)で説明してます。\nPyTorchの入門は[こちら](https://qiita.com/birdwatcher/items/e8ab9f6bba558759c106)\n\n参考にしたNotebook：\n- https://www.kaggle.com/konumaru/saint-with-pytorch-transformer-module\n- https://www.kaggle.com/manikanthr5/riiid-sakt-model-training-public","metadata":{}},{"cell_type":"markdown","source":"## データ読み込み","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom pathlib import Path\nimport time","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"input_dir = Path(\"/kaggle/input/riiid-test-answer-prediction/\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\ndata = pd.read_csv(input_dir / \"train.csv\",\n                   usecols=[1, 2, 3, 4, 7],\n                   dtype={\n                        'timestamp': 'int64',\n                        'user_id': 'int32',\n                        'content_id': 'int16',\n                        'content_type_id': 'int8',\n                        'answered_correctly':'int8',\n                   }\n       )\n# 2分くらいかかる","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 問題のマスタ\nquestions = pd.read_csv(input_dir/\"questions.csv\",\n                        dtype={\n                            \"question_id\":np.int16,\n                            \"bundle_id\":np.int16,\n                            \"correct_answer\":np.int8,\n                            \"part\":np.int8,\n                        }\n           )\nn_questions = len(questions)\nn_questions","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del questions","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### データ概要\n- 問題に正解したかどうか answerd_correctly $\\in \\{0,1\\}$ を予測するタスク\n- content_id: 問題のID/講義のID\n- content_type_id: 問題=0, 講義=1","metadata":{}},{"cell_type":"code","source":"# 講義のデータを落とす & 時系列順にソート\ndata = data[data[\"content_type_id\"] == 0].sort_values('timestamp').reset_index(drop = True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data.shape","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# メモリが貧弱なのでデータを絞る\ndata=data.tail(10000000)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n# ユーザーごとに「問題IDの系列」と「正解不正解の系列」を持つ形に変形\ndata = data.groupby(\"user_id\").apply(\n    lambda row: (\n        row[\"content_id\"].values,\n        row[\"answered_correctly\"].values,\n    )\n)\ndata","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# モデルで扱う最大系列長\nMAX_SEQ = 20\n# モデルで考慮する最小サンプル数（このサンプル数未満のユーザーは無視される）\nMIN_SAMPLES = 5\n# 埋め込み次元数\nEMBED_DIM = 32\n# Attentionヘッドの数\nNUM_HEADS = 2\n# ドロップアウト割合\nDROPOUT_RATE = 0.2\n# 学習率\nLEARNING_RATE = 1e-3\n# 最大学習率\nMAX_LEARNING_RATE = 2e-3\n# エポック数（学習データ全体を何周するか）\nEPOCHS = 5\n# バッチサイズ\nBATCH_SIZE = 1024","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 学習に使うユーザー数\nTRAIN_SAMPLES = int(data.shape[0]*0.8)\n# 学習データと検証データに分ける\ntrain_index = data.index.to_list()[:TRAIN_SAMPLES]\nvalid_index = data.index.to_list()[TRAIN_SAMPLES:]\ntrain = data[data.index.isin(train_index)]\nvalid = data[data.index.isin(valid_index)]\nprint(len(train), len(valid))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del data, train_index, valid_index","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## モデル作成","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.metrics import roc_auc_score\nfrom tqdm.notebook import tqdm","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 普通のTransfomer\n- いよいよ本題\n- まずは普通に思いつきそうなTransformerから\n\n![](https://camo.qiitausercontent.com/4d79e3addad65aa484d0cf9625de272187dc3b39/68747470733a2f2f71696974612d696d6167652d73746f72652e73332e61702d6e6f727468656173742d312e616d617a6f6e6177732e636f6d2f302f3239343636312f32373266663064382d353537352d333462372d366339372d3138363131386631343935382e706e67)","metadata":{}},{"cell_type":"markdown","source":"#### データローダー\n作るモデルに合ったデータセットの形式にする","metadata":{}},{"cell_type":"code","source":"# 系列の開始記号（Begin Of Sequence）\nBOS = 2\nclass TransformerDataset(Dataset):\n    \"\"\"普通の場合\"\"\"\n    def __init__(self, group, n_questions, min_samples=1, max_seq=128):\n        super(TransformerDataset, self).__init__()\n        self.max_seq = max_seq\n        self.n_questions = n_questions\n        # ユーザーID→系列を格納する変数\n        self.samples = {}\n        \n        self.user_ids = []\n        for user_id in group.index:\n            q, qa = group[user_id]\n            # サンプルが少ないユーザーは無視\n            if len(q) < min_samples:\n                continue\n            \n            # 最大系列長より長い系列の場合\n            if len(q) > self.max_seq:\n                total_questions = len(q)\n                # 最初の端数分の系列を格納\n                initial = total_questions % self.max_seq\n                if initial >= min_samples:\n                    self.user_ids.append(f\"{user_id}_0\")\n                    self.samples[f\"{user_id}_0\"] = (q[:initial], qa[:initial])\n                # 残りの長い系列について最大系列長ずつ取り出して格納\n                for seq in range(total_questions // self.max_seq):\n                    self.user_ids.append(f\"{user_id}_{seq+1}\")\n                    start = initial + seq * self.max_seq\n                    end = start + self.max_seq\n                    self.samples[f\"{user_id}_{seq+1}\"] = (q[start:end], qa[start:end])\n            else:\n                # 最大系列長より短い系列の場合\n                user_id = str(user_id)\n                self.user_ids.append(user_id)\n                self.samples[user_id] = (q, qa)\n    \n    def __len__(self):\n        return len(self.user_ids)\n\n    def __getitem__(self, index):\n        user_id = self.user_ids[index]\n        q_, qa_ = self.samples[user_id]\n        seq_len = len(q_)\n        # 最大系列長で揃える\n        q = np.zeros(self.max_seq, dtype=int)\n        qa = np.zeros(self.max_seq, dtype=int)\n        qa_shift = np.zeros(self.max_seq, dtype=int)\n        if seq_len == self.max_seq:\n            q[:] = q_\n            qa[:] = qa_\n        else:# 最大長ないものは末尾に格納\n            q[-seq_len:] = q_\n            qa[-seq_len:] = qa_\n        # 右シフトしたもの (Decoderの入力)\n        qa_shift[-seq_len:] = np.concatenate([[BOS],qa_[:-1]])\n        \n        # transformerで無視する部分を指定するマスク\n        # True: マスクされる、False: マスクなし\n        padding_mask = np.ones(self.max_seq, dtype=bool)\n        # 末尾のデータ格納箇所はマスクしない\n        padding_mask[-seq_len:] = False\n        \n        return q, qa_shift, qa, padding_mask","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset = TransformerDataset(train, n_questions, min_samples=MIN_SAMPLES, max_seq=MAX_SEQ)\ntrain_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=3)\nvalid_dataset = TransformerDataset(valid, n_questions, max_seq=MAX_SEQ)\nvalid_dataloader = DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=3)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Transformer","metadata":{}},{"cell_type":"code","source":"class FFN(nn.Module):\n    \"\"\"Feed Foward Network\"\"\"\n    def __init__(self, state_size=200):\n        super(FFN, self).__init__()\n        self.state_size = state_size\n        # 線形変換とReLU\n        self.lr1 = nn.Linear(state_size, state_size)\n        self.relu = nn.ReLU()\n        self.lr2 = nn.Linear(state_size, state_size)\n        self.dropout = nn.Dropout(0.2)\n    \n    def forward(self, x):\n        x = self.lr1(x)\n        x = self.relu(x)\n        x = self.lr2(x)\n        return self.dropout(x)\n\ndef future_mask(seq_length):\n    \"\"\"未来情報を見ないためのマスク\"\"\"\n    # 対角が0で右上の三角部分が1の上三角行列\n    future_mask = np.triu(np.ones((seq_length, seq_length)), k=1).astype('bool')\n    return torch.from_numpy(future_mask)\n\nclass Transformer(nn.Module):\n    def __init__(self, n_questions, n_response, max_seq=128, embed_dim=128, num_heads=8, dropout_rate=0.2):\n        super(Transformer, self).__init__()\n        self.n_questions = n_questions # 問題の数\n        self.n_response = n_response # {0, 1, BOS}の3種類\n        self.embed_dim = embed_dim\n        self.max_seq = max_seq\n        # Embedding系\n        # 離散表現（IDなどの整数）を指定した次元数の分散表現に変換する層（自然言語でいうところのword2vec）\n        # nn.Embedding(単語の種類数, 埋め込みたい次元数)\n        # nn.Embedding内でやっていることは単純で、onehotベクトルを作ってから線形変換しているだけっぽい\n        self.position_embed_e = nn.Embedding(max_seq, embed_dim)\n        self.question_embed_e = nn.Embedding(n_questions, embed_dim)\n        self.position_embed_d = nn.Embedding(max_seq, embed_dim)\n        self.response_embed_d = nn.Embedding(n_response, embed_dim)\n        # Attention系\n        # 直前のLinearはこの中に含まれている\n        # See https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html\n        self.attention_e = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, dropout=dropout_rate)\n        self.attention_d = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, dropout=dropout_rate)\n        self.attention_ed = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, dropout=dropout_rate)\n        # 正規化系\n        self.layer_normal1_e = nn.LayerNorm(embed_dim)\n        self.layer_normal2_e = nn.LayerNorm(embed_dim)\n        self.layer_normal1_d = nn.LayerNorm(embed_dim)\n        self.layer_normal2_d = nn.LayerNorm(embed_dim)\n        self.layer_normal3_d = nn.LayerNorm(embed_dim)\n        # FFN系\n        self.ffn_e = FFN(embed_dim)\n        self.ffn_d = FFN(embed_dim)\n        # 最後の線形変換\n        self.pred = nn.Linear(embed_dim, 1)\n    \n    def forward(self, question_ids, responses):\n        device = question_ids.device\n        #######################################################################\n        # Encoder\n        #######################################################################\n        # --------- Embedding ---------\n        # 系列の何番目かを表す[[0, 1, 2, ..., seq_len - 1]]を生成\n        seq = torch.arange(self.max_seq, device=device).unsqueeze(0)\n        # それをEmbedding（分散表現へ）\n        pos_e = self.position_embed_e(seq)\n        # Encoderに入れるID列を分散表現へ\n        question = self.question_embed_e(question_ids)\n        # 位置情報を加える\n        enc = pos_e + question\n        # --------- Attention ---------\n        # Attentionの関数は、（系列長, バッチサイズ, 次元数）の形で受け取る\n        enc = enc.permute(1, 0, 2) # （バッチサイズ, 系列長, 次元数）=>（系列長, バッチサイズ, 次元数）\n        # self-attention\n        enc_tmp, _ = self.attention_e(enc, enc, enc)\n        enc = self.layer_normal1_e(enc_tmp+enc)\n        # --------- Feed Forward Network ---------\n        enc_tmp = self.ffn_e(enc)\n        enc = self.layer_normal2_e(enc_tmp+enc)\n        #######################################################################\n        # Decoder\n        #######################################################################\n        # --------- Embedding ---------\n        # 位置情報をEmbedding（分散表現へ）\n        pos_d = self.position_embed_d(seq)\n        # Decoderに入れる離散表現を分散表現へ\n        response = self.response_embed_d(responses)\n        dec = response + pos_d\n        # --------- Attention ---------\n        # 未来情報を隠すマスク\n        att_mask = future_mask(self.max_seq).to(device)\n        # Attentionの関数は、（系列長, バッチサイズ, 次元数）の形で受け取る\n        dec = dec.permute(1, 0, 2) # （バッチサイズ, 系列長, 次元数）=>（系列長, バッチサイズ, 次元数）\n        # self-attention\n        dec_tmp, _ = self.attention_d(dec, dec, dec, attn_mask=att_mask)\n        dec = self.layer_normal1_d(dec_tmp+dec)\n        # source-target attention\n        dec_tmp, _ = self.attention_ed(dec, enc, enc)\n        dec = self.layer_normal2_d(dec_tmp+dec)\n        # --------- Feed Forward Network ---------\n        dec_tmp = self.ffn_d(dec)\n        dec = self.layer_normal3_d(dec_tmp+dec)\n        # shapeをもとに戻す\n        dec = dec.permute(1, 0, 2) # （系列長, バッチサイズ, 次元数）=>（バッチサイズ, 系列長, 次元数）\n\n        # 最後の線形変換\n        dec = self.pred(dec)\n\n        return dec.squeeze(-1)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"- 今回は[nn.MultiheadAttention](https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html)を使って他はすべて自分で書いた\n- AttentionからFeed Forward Networkまでの一連の内容をサボりたければ[nn.Transformer](https://pytorch.org/docs/stable/generated/torch.nn.Transformer.html)モジュールを使う方法もある\n    - [実装を確認](https://pytorch.org/docs/stable/_modules/torch/nn/modules/transformer.html#Transformer)すると、Attention, FFN, 正規化までやっているようだ\n    - これを使えば、自分で書くべき箇所はEmbeddingと最後の線形変換だけになる\n- Attentionの実装が気になる場合は[こちら](https://github.com/jadore801120/attention-is-all-you-need-pytorch/tree/132907dd272e2cc92e3c10e6c4e783a87ff8893d)の実装がわかりやすいと感じた\n    - ライブラリにあるnn.MultiheadAttentionの実装ではなく、独自でMultiHeadAttentionを書いている","metadata":{}},{"cell_type":"markdown","source":"#### 最適化法と損失の設定","metadata":{}},{"cell_type":"code","source":"# 0: 不正解、1：正解、2：BOS\nn_responses = 3\n\n# モデルの設定\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nmodel = Transformer(n_questions, n_responses,\n                    max_seq=MAX_SEQ,\n                    embed_dim=EMBED_DIM,\n                    num_heads=NUM_HEADS,\n                   )\n# 最適化法の指定\noptimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)\n# 損失の設定\n# - BCEWithLogitsLoss: sigmoidをとる前の値を入力として受け付ける\n# - BCELoss: sigmoidをとった後の値を入力として受け付ける\n# モデル自体にsigmoidを含めてないため、前者を使う（この方が数値的に安定していて、高速らしい）\ncriterion = nn.BCEWithLogitsLoss()\n# 学習率をどう動かしていくか\nscheduler = torch.optim.lr_scheduler.OneCycleLR(\n    optimizer, max_lr=MAX_LEARNING_RATE, steps_per_epoch=len(train_dataloader), epochs=EPOCHS\n)\n\nmodel.to(device)\ncriterion.to(device)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### 学習と評価のループ","metadata":{}},{"cell_type":"code","source":"def train_fn(model, dataloader, optimizer, scheduler, criterion, device=\"cpu\"):\n    \"\"\"学習ループ\"\"\"\n    model.train()\n\n    train_loss = []\n    num_corrects = 0\n    num_total = 0\n    targets = []\n    outs = []\n\n    for item in tqdm(dataloader):\n        # tensorを置くデバイスの設定\n        enc_input = item[0].to(device).long()\n        dec_input = item[1].to(device).long()\n        target = item[2].to(device).float()\n        padding_mask = item[3].to(device).bool()\n        # 勾配情報の初期化\n        optimizer.zero_grad()\n        # モデルの出力計算\n        output = model(enc_input, dec_input)\n        # 予測値のうち、maskされていない場所だけ取り出す\n        output = torch.masked_select(output, torch.logical_not(padding_mask))\n        # 真値のうち、maskされていない場所だけ取り出す\n        target = torch.masked_select(target, torch.logical_not(padding_mask))\n        # 損失を計算\n        loss = criterion(output, target)\n        # 勾配を計算\n        loss.backward()\n        # 最適化法に基づいてパラメータ更新\n        optimizer.step()\n        scheduler.step()\n        # 損失を記録\n        train_loss.append(loss.item())\n        # 0.5以上なら1と予測\n        pred = (torch.sigmoid(output) >= 0.5).long()\n        # 精度のための記録\n        num_corrects += (pred == target).sum().item()\n        num_total += len(target)\n\n        targets.extend(target.view(-1).data.cpu().numpy())\n        outs.extend(output.view(-1).data.cpu().numpy())\n\n    acc = num_corrects / num_total\n    auc = roc_auc_score(targets, outs)\n    loss = np.mean(train_loss)\n\n    return loss, acc, auc","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def valid_fn(model, dataloader, criterion, device=\"cpu\"):\n    \"\"\"検証のループ\"\"\"\n    model.eval()\n\n    valid_loss = []\n    num_corrects = 0\n    num_total = 0\n    targets = []\n    outs = []\n\n    for item in tqdm(dataloader):\n        enc_input = item[0].to(device).long()\n        dec_input = item[1].to(device).long()\n        target = item[2].to(device).float()\n        padding_mask = item[3].to(device).bool()\n        # モデルの出力計算\n        output = model(enc_input, dec_input)\n        # 予測値のうち、maskされていない場所だけ取り出す\n        output = torch.masked_select(output, torch.logical_not(padding_mask))\n        # 真値のうち、maskされていない場所だけ取り出す\n        target = torch.masked_select(target, torch.logical_not(padding_mask))\n        # 損失を計算\n        loss = criterion(output, target)\n        # 損失を記録\n        valid_loss.append(loss.item())\n        # 0.5以上なら1と予測\n        pred = (torch.sigmoid(output) >= 0.5).long()\n        # 精度のための記録\n        num_corrects += (pred == target).sum().item()\n        num_total += len(target)\n\n        targets.extend(target.view(-1).data.cpu().numpy())\n        outs.extend(output.view(-1).data.cpu().numpy())\n\n    acc = num_corrects / num_total\n    auc = roc_auc_score(targets, outs)\n    loss = np.mean(valid_loss)\n\n    return loss, acc, auc","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"best_auc = 0\nearly_stop = 3\nstep = 0\nfor epoch in range(EPOCHS):\n    loss, acc, auc = train_fn(model, train_dataloader, optimizer, scheduler, criterion, device)\n    print(f\"[train] epoch: {epoch+1}/{EPOCHS}, loss: {loss:.3f}, acc: {acc:.3f}, auc: {auc:.3f}\")\n    loss, acc, auc = valid_fn(model, valid_dataloader, criterion, device)\n    print(f\"[valid] epoch: {epoch+1}/{EPOCHS}, loss: {loss:.3f}, acc: {acc:.3f}, auc: {auc:.3f}\")\n    if auc > best_auc:\n        best_auc = auc\n        step = 0\n        torch.save(model.state_dict(), \"model.pt\")\n    else:\n        step += 1\n        if step >= early_stop:\n            break","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### SAKT model\n- [A Self-Attentive model for Knowledge Tracing](https://arxiv.org/pdf/1907.06837.pdf)というものがあるらしい\n- Encoder/Decoderモデルではなく, 「問題IDと正解不正解の系列をペア」と「問題ID」をAttentionするモデル\n\n![](https://camo.qiitausercontent.com/a99c49a8181f70ef45f8241a70e42245eec3197b/68747470733a2f2f71696974612d696d6167652d73746f72652e73332e61702d6e6f727468656173742d312e616d617a6f6e6177732e636f6d2f302f3239343636312f33383134376631352d663164352d663061652d333139372d6531613966616266373836632e706e67)","metadata":{}},{"cell_type":"markdown","source":"#### データローダー\n作るモデルに合ったデータセットの形式にする","metadata":{}},{"cell_type":"code","source":"class SAKTDataset(Dataset):\n    def __init__(self, group, n_questions, min_samples=1, max_seq=128):\n        super(SAKTDataset, self).__init__()\n        self.max_seq = max_seq\n        self.n_questions = n_questions\n        # ユーザーID→系列を格納する変数\n        self.samples = {}\n        \n        self.user_ids = []\n        for user_id in group.index:\n            q, qa = group[user_id]\n            # サンプルが少ないユーザーは無視\n            if len(q) < min_samples:\n                continue\n            \n            # 最大系列長より長い系列の場合\n            if len(q) > self.max_seq:\n                total_questions = len(q)\n                # 最初の端数分の系列を格納\n                initial = total_questions % self.max_seq\n                if initial >= min_samples:\n                    self.user_ids.append(f\"{user_id}_0\")\n                    self.samples[f\"{user_id}_0\"] = (q[:initial], qa[:initial])\n                # 残りの長い系列について最大系列長ずつ取り出して格納\n                for seq in range(total_questions // self.max_seq):\n                    self.user_ids.append(f\"{user_id}_{seq+1}\")\n                    start = initial + seq * self.max_seq\n                    end = start + self.max_seq\n                    self.samples[f\"{user_id}_{seq+1}\"] = (q[start:end], qa[start:end])\n            else:\n                # 最大系列長より短い系列の場合\n                user_id = str(user_id)\n                self.user_ids.append(user_id)\n                self.samples[user_id] = (q, qa)\n    \n    def __len__(self):\n        return len(self.user_ids)\n\n    def __getitem__(self, index):\n        user_id = self.user_ids[index]\n        q_, qa_ = self.samples[user_id]\n        seq_len = len(q_)\n        # 最大系列長で揃える\n        q = np.zeros(self.max_seq, dtype=int)\n        qa = np.zeros(self.max_seq, dtype=int)\n        if seq_len == self.max_seq:\n            q[:] = q_\n            qa[:] = qa_\n        else:\n            q[-seq_len:] = q_\n            qa[-seq_len:] = qa_\n        # 右シフトした問題のID系列 \n        target_id = q[1:]\n        # 正解不正解の系列 (予測したいもの)\n        label = qa[1:]\n\n        # 正解：コンテンツID + 全質問数\n        # 不正解：コンテンツID\n        # 数字で列挙すると、不正解が並んだあとに、正解が並ぶイメージ\n        # そのような数字が並んだ系列を入力にする\n        # 最後の要素は予測対象なため除外\n        x = q[:-1] + (qa[:-1] == 1) * self.n_questions\n\n        # 無視する部分を指定するマスク\n        # True: マスクされる、False: マスクなし\n        padding_mask = np.ones(self.max_seq, dtype=bool)\n        # 末尾のデータ格納箇所はマスクしない\n        padding_mask[-seq_len:] = False\n        padding_mask = padding_mask[1:]\n\n        return x, target_id, label, padding_mask","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset = SAKTDataset(train, n_questions, min_samples=MIN_SAMPLES, max_seq=MAX_SEQ)\ntrain_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=3)\nvalid_dataset = SAKTDataset(valid, n_questions, max_seq=MAX_SEQ)\nvalid_dataloader = DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=3)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### SAKT","metadata":{}},{"cell_type":"code","source":"class SAKT(nn.Module):\n    \"\"\"SAKTモデル\"\"\"\n    def __init__(self, n_questions, max_seq=128, embed_dim=128, num_heads=8, dropout_rate=0.2):\n        super(SAKT, self).__init__()\n        self.n_questions = n_questions\n        self.embed_dim = embed_dim\n\n        self.embedding = nn.Embedding(2*n_questions, embed_dim)\n        self.pos_embedding = nn.Embedding(max_seq-1, embed_dim)\n        self.e_embedding = nn.Embedding(n_questions, embed_dim)\n        # MultiheadAttentionというクラスが用意されている\n        self.multi_att = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, dropout=dropout_rate)\n\n        self.dropout = nn.Dropout(dropout_rate)\n        self.layer_normal = nn.LayerNorm(embed_dim) \n\n        self.ffn = FFN(embed_dim)\n        self.pred = nn.Linear(embed_dim, 1)\n    \n    def forward(self, x, question_ids):\n        device = x.device\n        # ID系列の埋め込み\n        x = self.embedding(x)\n        # 系列の何番目かを表す[[0, 1, 2, ..., seq_len - 1]]を生成\n        pos_id = torch.arange(x.size(1)).unsqueeze(0).to(device)\n        # 系列内位置情報の埋め込み\n        pos_x = self.pos_embedding(pos_id)\n\n        x = x + pos_x\n        # 問題の系列も埋め込み\n        e = self.e_embedding(question_ids)\n\n        x = x.permute(1, 0, 2) # x: [bs, s_len, embed] => [s_len, bs, embed]\n        e = e.permute(1, 0, 2)\n        # 未来情報を隠すマスク\n        att_mask = future_mask(x.size(0)).to(device)\n        att_output, att_weight = self.multi_att(e, x, x, attn_mask=att_mask)\n        att_output = self.layer_normal(att_output + e)\n        att_output = att_output.permute(1, 0, 2) # att_output: [s_len, bs, embed] => [bs, s_len, embed]\n\n        x = self.ffn(att_output)\n        x = self.layer_normal(x + att_output)\n        x = self.pred(x)\n\n        return x.squeeze(-1)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"- `FNN`と`future_mask`は、普通のTransformerの方で定義したものを使ってます","metadata":{}},{"cell_type":"markdown","source":"#### 最適化法と損失の設定","metadata":{}},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nmodel = SAKT(n_questions, max_seq=MAX_SEQ, embed_dim=EMBED_DIM, num_heads=NUM_HEADS, dropout_rate=DROPOUT_RATE)\noptimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)\ncriterion = nn.BCEWithLogitsLoss()\nscheduler = torch.optim.lr_scheduler.OneCycleLR(\n    optimizer, max_lr=MAX_LEARNING_RATE, steps_per_epoch=len(train_dataloader), epochs=EPOCHS\n)\n\nmodel.to(device)\ncriterion.to(device)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### 学習\n学習のループは普通のTransformerの方の関数と同じものを使います。","metadata":{}},{"cell_type":"code","source":"best_auc = 0\nearly_stop = 3\nstep = 0\nfor epoch in range(EPOCHS):\n    loss, acc, auc = train_fn(model, train_dataloader, optimizer, scheduler, criterion, device)\n    print(f\"[train] epoch: {epoch+1}/{EPOCHS}, loss: {loss:.3f}, acc: {acc:.3f}, auc: {auc:.3f}\")\n    loss, acc, auc = valid_fn(model, valid_dataloader, criterion, device)\n    print(f\"[valid] epoch: {epoch+1}/{EPOCHS}, loss: {loss:.3f}, acc: {acc:.3f}, auc: {auc:.3f}\")\n    if auc > best_auc:\n        best_auc = auc\n        step = 0\n        torch.save(model.state_dict(), \"sakt_model.pt\")\n    else:\n        step += 1\n        if step >= early_stop:\n            break","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}