20230330 self attention pooling public 0.95750

ML2023_hw04.ipynb

@@ -263,6 +263,11 @@
     "        \n",
     "        # Project the the dimension of features from d_model into speaker nums.\n",
     "        self.pred_layer = nn.Linear(d_model, n_spks)\n",
+    "        \n",
+    "        # softmax\n",
+    "        self.softmax = nn.functional.softmax\n",
+    "        \n",
+    "        self.W = nn.Linear(d_model, 1)\n",
     "\n",
     "    def forward(self, mels):\n",
     "        \"\"\"\n",
@@ -280,15 +285,18 @@
     "        #lengths = torch.randint(out.size(1), (int(out.size(0)),)).to(device)\n",
     "        #lengths[torch.argmax(lengths)] = out.size(1)\n",
     "        lengths = torch.full((int(out.size(0)),), out.size(1)).to(device)\n",
-    "        out, _ = self.encoder(out, lengths)\n",
     "        # out: (batch size, length, d_model)\n",
-    "        #out = out.transpose(0, 1)\n",
+    "        out, _ = self.encoder(out, lengths)\n",
     "        # mean pooling\n",
-    "        #stats = out.mean(dim=1)\n",
-    "\n",
-    "        # out: (batch, n_spks)\n",
     "        stats = out.mean(dim=1)\n",
-    "        out = self.pred_layer(stats)\n",
+    "        # self attention pooling \n",
+    "        # reference from https://gist.github.com/pohanchi/c77f6dbfbcbc21c5215acde4f62e4362\n",
+    "        # input batch_rep : size (N, T, H), N: batch size, T: sequence length, H: Hidden dimension\n",
+    "        # output utter_rep: size (N, H)\n",
+    "        att_w = self.softmax(self.W(out).squeeze(-1), dim=1).unsqueeze(-1)\n",
+    "        utter_rep = torch.sum(out * att_w, dim=1)\n",
+    "        \n",
+    "        out = self.pred_layer(utter_rep)\n",
     "        return out, _"
    ]
   },
@@ -308,7 +316,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 7,
    "metadata": {
     "id": "ykt0N1nVJJi2"
    },
@@ -376,7 +384,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 8,
    "metadata": {
     "id": "N-rr8529JMz0"
    },
@@ -416,7 +424,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 9,
    "metadata": {
     "id": "YAiv6kpdJRTJ"
    },
@@ -485,16 +493,280 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Train: 100% 2000/2000 [04:36<00:00,  7.24 step/s, accuracy=0.19, loss=3.68, step=2000]  \n",
-      "Valid: 100% 5664/5667 [00:52<00:00, 107.38 uttr/s, accuracy=0.23, loss=3.66]\n",
-      "Train: 100% 2000/2000 [04:33<00:00,  7.30 step/s, accuracy=0.47, loss=2.42, step=4000]  \n",
-      "Valid: 100% 5664/5667 [00:55<00:00, 101.57 uttr/s, accuracy=0.40, loss=2.75]\n",
-      "Train: 100% 2000/2000 [04:53<00:00,  6.82 step/s, accuracy=0.53, loss=1.94, step=6000] \n",
-      "Valid: 100% 5664/5667 [00:54<00:00, 103.31 uttr/s, accuracy=0.51, loss=2.22]\n",
-      "Train: 100% 2000/2000 [05:36<00:00,  5.94 step/s, accuracy=0.69, loss=1.35, step=8000] \n",
-      "Valid: 100% 5664/5667 [00:52<00:00, 108.07 uttr/s, accuracy=0.58, loss=1.88]\n",
-      "Train: 100% 2000/2000 [04:40<00:00,  7.13 step/s, accuracy=0.75, loss=0.93, step=1e+4]  \n",
-      "Valid: 100% 5664/5667 [00:53<00:00, 106.50 uttr/s, accuracy=0.63, loss=1.69]\n",
+      "Train: 100% 2000/2000 [04:31<00:00,  7.36 step/s, accuracy=0.22, loss=3.40, step=2000]  \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 106.62 uttr/s, accuracy=0.24, loss=3.57]\n",
+      "Train: 100% 2000/2000 [04:38<00:00,  7.17 step/s, accuracy=0.34, loss=2.71, step=4000]  \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 104.90 uttr/s, accuracy=0.42, loss=2.59]\n",
+      "Train: 100% 2000/2000 [04:33<00:00,  7.31 step/s, accuracy=0.59, loss=1.64, step=6000] \n",
+      "Valid: 100% 5664/5667 [00:54<00:00, 103.97 uttr/s, accuracy=0.55, loss=2.03]\n",
+      "Train: 100% 2000/2000 [12:48<00:00,  2.60 step/s, accuracy=0.56, loss=2.18, step=8000]   \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 108.64 uttr/s, accuracy=0.60, loss=1.74]\n",
+      "Train: 100% 2000/2000 [04:35<00:00,  7.26 step/s, accuracy=0.72, loss=1.23, step=1e+4]  \n",
+      "Valid: 100% 5664/5667 [00:54<00:00, 103.98 uttr/s, accuracy=0.64, loss=1.60]\n",
+      "Train:   0% 0/2000 [00:00<?, ? step/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 10000, best model saved. (accuracy=0.6398)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Train: 100% 2000/2000 [04:35<00:00,  7.26 step/s, accuracy=0.72, loss=1.43, step=12000]  \n",
+      "Valid: 100% 5664/5667 [00:54<00:00, 103.49 uttr/s, accuracy=0.66, loss=1.55]\n",
+      "Train: 100% 2000/2000 [04:37<00:00,  7.22 step/s, accuracy=0.88, loss=0.73, step=14000] \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 105.48 uttr/s, accuracy=0.70, loss=1.34]\n",
+      "Train: 100% 2000/2000 [05:48<00:00,  5.74 step/s, accuracy=0.72, loss=1.05, step=16000] \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 108.00 uttr/s, accuracy=0.72, loss=1.28]\n",
+      "Train: 100% 2000/2000 [04:54<00:00,  6.79 step/s, accuracy=0.84, loss=0.52, step=18000]  \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 107.34 uttr/s, accuracy=0.73, loss=1.28]\n",
+      "Train: 100% 2000/2000 [05:03<00:00,  6.59 step/s, accuracy=0.81, loss=0.69, step=2e+4]   \n",
+      "Valid: 100% 5664/5667 [00:54<00:00, 103.80 uttr/s, accuracy=0.73, loss=1.25]\n",
+      "Train:   0% 0/2000 [00:00<?, ? step/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 20000, best model saved. (accuracy=0.7334)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Train: 100% 2000/2000 [05:20<00:00,  6.24 step/s, accuracy=0.91, loss=0.48, step=22000] \n",
+      "Valid: 100% 5664/5667 [00:56<00:00, 99.70 uttr/s, accuracy=0.74, loss=1.22] \n",
+      "Train: 100% 2000/2000 [05:47<00:00,  5.75 step/s, accuracy=0.78, loss=0.56, step=24000] \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 108.01 uttr/s, accuracy=0.75, loss=1.21]\n",
+      "Train: 100% 2000/2000 [04:58<00:00,  6.69 step/s, accuracy=0.88, loss=0.34, step=26000]  \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 105.23 uttr/s, accuracy=0.75, loss=1.20]\n",
+      "Train: 100% 2000/2000 [05:29<00:00,  6.06 step/s, accuracy=0.88, loss=0.37, step=28000]  \n",
+      "Valid: 100% 5664/5667 [00:56<00:00, 100.46 uttr/s, accuracy=0.77, loss=1.12]\n",
+      "Train: 100% 2000/2000 [05:06<00:00,  6.52 step/s, accuracy=0.88, loss=0.58, step=3e+4]  \n",
+      "Valid: 100% 5664/5667 [00:55<00:00, 101.50 uttr/s, accuracy=0.78, loss=1.05]\n",
+      "Train:   0% 0/2000 [00:00<?, ? step/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 30000, best model saved. (accuracy=0.7800)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Train: 100% 2000/2000 [05:52<00:00,  5.67 step/s, accuracy=0.94, loss=0.20, step=32000] \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 108.10 uttr/s, accuracy=0.79, loss=1.06]\n",
+      "Train: 100% 2000/2000 [04:51<00:00,  6.87 step/s, accuracy=0.84, loss=0.73, step=34000]  \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 106.68 uttr/s, accuracy=0.78, loss=1.09]\n",
+      "Train: 100% 2000/2000 [04:43<00:00,  7.05 step/s, accuracy=0.88, loss=0.55, step=36000]  \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 105.69 uttr/s, accuracy=0.78, loss=1.08]\n",
+      "Train: 100% 2000/2000 [04:47<00:00,  6.97 step/s, accuracy=0.91, loss=0.39, step=38000] \n",
+      "Valid: 100% 5664/5667 [00:55<00:00, 102.95 uttr/s, accuracy=0.79, loss=1.06]\n",
+      "Train: 100% 2000/2000 [05:34<00:00,  5.98 step/s, accuracy=0.91, loss=0.51, step=4e+4]  \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 108.32 uttr/s, accuracy=0.81, loss=0.97]\n",
+      "Train:   0% 0/2000 [00:00<?, ? step/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 40000, best model saved. (accuracy=0.8069)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Train: 100% 2000/2000 [04:44<00:00,  7.04 step/s, accuracy=0.97, loss=0.38, step=42000]  \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 105.66 uttr/s, accuracy=0.79, loss=1.01]\n",
+      "Train: 100% 2000/2000 [05:02<00:00,  6.61 step/s, accuracy=0.84, loss=0.46, step=44000]  \n",
+      "Valid: 100% 5664/5667 [00:55<00:00, 102.54 uttr/s, accuracy=0.80, loss=1.01]\n",
+      "Train: 100% 2000/2000 [05:07<00:00,  6.51 step/s, accuracy=0.78, loss=0.68, step=46000] \n",
+      "Valid: 100% 5664/5667 [00:57<00:00, 97.74 uttr/s, accuracy=0.80, loss=0.99] \n",
+      "Train: 100% 2000/2000 [05:44<00:00,  5.80 step/s, accuracy=0.94, loss=0.35, step=48000] \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 107.12 uttr/s, accuracy=0.81, loss=0.98]\n",
+      "Train: 100% 2000/2000 [04:51<00:00,  6.85 step/s, accuracy=0.88, loss=0.30, step=5e+4]   \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 105.84 uttr/s, accuracy=0.81, loss=0.98]\n",
+      "Train:   0% 0/2000 [00:00<?, ? step/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 50000, best model saved. (accuracy=0.8120)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Train: 100% 2000/2000 [05:03<00:00,  6.59 step/s, accuracy=0.81, loss=0.74, step=52000] \n",
+      "Valid: 100% 5664/5667 [00:54<00:00, 104.29 uttr/s, accuracy=0.81, loss=0.93]\n",
+      "Train: 100% 2000/2000 [05:12<00:00,  6.40 step/s, accuracy=0.78, loss=0.70, step=54000] \n",
+      "Valid: 100% 5664/5667 [00:56<00:00, 100.11 uttr/s, accuracy=0.82, loss=0.95]\n",
+      "Train: 100% 2000/2000 [05:59<00:00,  5.56 step/s, accuracy=0.94, loss=0.25, step=56000] \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 107.08 uttr/s, accuracy=0.82, loss=0.94]\n",
+      "Train: 100% 2000/2000 [05:04<00:00,  6.58 step/s, accuracy=0.94, loss=0.16, step=58000]  \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 106.15 uttr/s, accuracy=0.81, loss=0.93]\n",
+      "Train: 100% 2000/2000 [05:08<00:00,  6.48 step/s, accuracy=0.78, loss=0.60, step=6e+4]  \n",
+      "Valid: 100% 5664/5667 [00:55<00:00, 101.84 uttr/s, accuracy=0.82, loss=0.91]\n",
+      "Train:   0% 0/2000 [00:00<?, ? step/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 60000, best model saved. (accuracy=0.8245)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Train: 100% 2000/2000 [05:16<00:00,  6.33 step/s, accuracy=0.91, loss=0.34, step=62000] \n",
+      "Valid: 100% 5664/5667 [00:55<00:00, 102.56 uttr/s, accuracy=0.83, loss=0.88]\n",
+      "Train: 100% 2000/2000 [18:59<00:00,  1.76 step/s, accuracy=0.97, loss=0.21, step=64000]   \n",
+      "Valid: 100% 5664/5667 [11:08<00:00,  8.48 uttr/s, accuracy=0.84, loss=0.87] \n",
+      "Train: 100% 2000/2000 [10:41<00:00,  3.12 step/s, accuracy=0.84, loss=0.54, step=66000]    \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 108.57 uttr/s, accuracy=0.83, loss=0.87]\n",
+      "Train: 100% 2000/2000 [04:28<00:00,  7.46 step/s, accuracy=0.91, loss=0.22, step=68000] \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 107.19 uttr/s, accuracy=0.83, loss=0.87]\n",
+      "Train: 100% 2000/2000 [04:25<00:00,  7.53 step/s, accuracy=0.97, loss=0.15, step=7e+4]  \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 106.26 uttr/s, accuracy=0.83, loss=0.89]\n",
+      "Train:   0% 0/2000 [00:00<?, ? step/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 70000, best model saved. (accuracy=0.8353)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Train: 100% 2000/2000 [10:44<00:00,  3.10 step/s, accuracy=0.97, loss=0.14, step=72000]  \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 107.85 uttr/s, accuracy=0.83, loss=0.90]\n",
+      "Train: 100% 2000/2000 [04:31<00:00,  7.37 step/s, accuracy=0.94, loss=0.22, step=74000]  \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 106.09 uttr/s, accuracy=0.83, loss=0.91]\n",
+      "Train: 100% 2000/2000 [04:31<00:00,  7.38 step/s, accuracy=0.88, loss=0.51, step=76000] \n",
+      "Valid: 100% 5664/5667 [00:55<00:00, 101.38 uttr/s, accuracy=0.83, loss=0.86]\n",
+      "Train: 100% 2000/2000 [04:52<00:00,  6.83 step/s, accuracy=0.84, loss=0.30, step=78000] \n",
+      "Valid: 100% 5664/5667 [00:54<00:00, 103.45 uttr/s, accuracy=0.84, loss=0.86]\n",
+      "Train: 100% 2000/2000 [05:30<00:00,  6.05 step/s, accuracy=0.91, loss=0.33, step=8e+4]   \n",
+      "Valid: 100% 5664/5667 [00:54<00:00, 104.53 uttr/s, accuracy=0.85, loss=0.86]\n",
+      "Train:   0% 0/2000 [00:00<?, ? step/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 80000, best model saved. (accuracy=0.8452)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Train: 100% 2000/2000 [04:54<00:00,  6.80 step/s, accuracy=0.94, loss=0.18, step=82000]  \n",
+      "Valid: 100% 5664/5667 [00:54<00:00, 103.49 uttr/s, accuracy=0.84, loss=0.84]\n",
+      "Train: 100% 2000/2000 [05:15<00:00,  6.35 step/s, accuracy=0.91, loss=0.36, step=84000] \n",
+      "Valid: 100% 5664/5667 [00:55<00:00, 101.81 uttr/s, accuracy=0.85, loss=0.82]\n",
+      "Train: 100% 2000/2000 [05:12<00:00,  6.41 step/s, accuracy=0.94, loss=0.25, step=86000] \n",
+      "Valid: 100% 5664/5667 [00:56<00:00, 99.54 uttr/s, accuracy=0.85, loss=0.84] \n",
+      "Train: 100% 2000/2000 [05:35<00:00,  5.96 step/s, accuracy=0.97, loss=0.12, step=88000]  \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 107.23 uttr/s, accuracy=0.86, loss=0.80]\n",
+      "Train: 100% 2000/2000 [05:00<00:00,  6.66 step/s, accuracy=0.97, loss=0.19, step=9e+4]   \n",
+      "Valid: 100% 5664/5667 [00:54<00:00, 104.47 uttr/s, accuracy=0.85, loss=0.85]\n",
+      "Train:   0% 0/2000 [00:00<?, ? step/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 90000, best model saved. (accuracy=0.8598)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Train: 100% 2000/2000 [04:59<00:00,  6.68 step/s, accuracy=0.94, loss=0.20, step=92000] \n",
+      "Valid: 100% 5664/5667 [00:55<00:00, 101.47 uttr/s, accuracy=0.85, loss=0.83]\n",
+      "Train: 100% 2000/2000 [06:05<00:00,  5.47 step/s, accuracy=0.91, loss=0.20, step=94000] \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 108.13 uttr/s, accuracy=0.86, loss=0.73]\n",
+      "Train: 100% 2000/2000 [05:17<00:00,  6.29 step/s, accuracy=1.00, loss=0.12, step=96000]  \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 105.82 uttr/s, accuracy=0.86, loss=0.80]\n",
+      "Train: 100% 2000/2000 [05:10<00:00,  6.45 step/s, accuracy=0.94, loss=0.23, step=98000]  \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 105.24 uttr/s, accuracy=0.86, loss=0.82]\n",
+      "Train: 100% 2000/2000 [05:13<00:00,  6.39 step/s, accuracy=0.94, loss=0.23, step=1e+5]  \n",
+      "Valid: 100% 5664/5667 [00:56<00:00, 101.03 uttr/s, accuracy=0.86, loss=0.77]\n",
+      "Train:   0% 0/2000 [00:00<?, ? step/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 100000, best model saved. (accuracy=0.8630)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Train: 100% 2000/2000 [05:58<00:00,  5.57 step/s, accuracy=0.88, loss=0.33, step=102000]\n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 106.35 uttr/s, accuracy=0.86, loss=0.75]\n",
+      "Train: 100% 2000/2000 [05:16<00:00,  6.33 step/s, accuracy=0.94, loss=0.21, step=104000]  \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 105.24 uttr/s, accuracy=0.86, loss=0.75]\n",
+      "Train: 100% 2000/2000 [05:14<00:00,  6.36 step/s, accuracy=0.97, loss=0.17, step=106000]  \n",
+      "Valid: 100% 5664/5667 [00:55<00:00, 102.57 uttr/s, accuracy=0.87, loss=0.75]\n",
+      "Train: 100% 2000/2000 [05:15<00:00,  6.34 step/s, accuracy=0.97, loss=0.07, step=108000] \n",
+      "Valid: 100% 5664/5667 [00:57<00:00, 99.33 uttr/s, accuracy=0.87, loss=0.73] \n",
+      "Train: 100% 2000/2000 [06:00<00:00,  5.55 step/s, accuracy=0.97, loss=0.11, step=110000] \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 105.74 uttr/s, accuracy=0.87, loss=0.73]\n",
+      "Train:   0% 0/2000 [00:00<?, ? step/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 110000, best model saved. (accuracy=0.8669)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Train: 100% 2000/2000 [05:11<00:00,  6.42 step/s, accuracy=1.00, loss=0.03, step=112000]  \n",
+      "Valid: 100% 5664/5667 [00:54<00:00, 104.34 uttr/s, accuracy=0.87, loss=0.72]\n",
+      "Train: 100% 2000/2000 [05:13<00:00,  6.37 step/s, accuracy=0.94, loss=0.13, step=114000]  \n",
+      "Valid: 100% 5664/5667 [00:55<00:00, 102.10 uttr/s, accuracy=0.87, loss=0.73]\n",
+      "Train: 100% 2000/2000 [05:28<00:00,  6.10 step/s, accuracy=0.91, loss=0.43, step=116000] \n",
+      "Valid: 100% 5664/5667 [00:56<00:00, 100.56 uttr/s, accuracy=0.87, loss=0.72]\n",
+      "Train: 100% 2000/2000 [06:08<00:00,  5.43 step/s, accuracy=1.00, loss=0.04, step=118000] \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 107.06 uttr/s, accuracy=0.88, loss=0.70]\n",
+      "Train: 100% 2000/2000 [05:18<00:00,  6.28 step/s, accuracy=1.00, loss=0.03, step=120000]  \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 105.12 uttr/s, accuracy=0.88, loss=0.73]\n",
       "Train:   0% 0/2000 [00:00<?, ? step/s]"
      ]
     },
@@ -502,23 +774,23 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Step 10000, best model saved. (accuracy=0.6262)\n"
+      "Step 120000, best model saved. (accuracy=0.8773)\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Train: 100% 2000/2000 [04:29<00:00,  7.43 step/s, accuracy=0.66, loss=1.40, step=12000]  \n",
-      "Valid: 100% 5664/5667 [00:54<00:00, 104.72 uttr/s, accuracy=0.65, loss=1.62]\n",
-      "Train: 100% 2000/2000 [04:29<00:00,  7.42 step/s, accuracy=0.62, loss=1.48, step=14000] \n",
-      "Valid: 100% 5664/5667 [00:54<00:00, 103.87 uttr/s, accuracy=0.67, loss=1.49]\n",
-      "Train: 100% 2000/2000 [05:17<00:00,  6.29 step/s, accuracy=0.69, loss=1.37, step=16000] \n",
-      "Valid: 100% 5664/5667 [00:52<00:00, 108.50 uttr/s, accuracy=0.68, loss=1.48] \n",
-      "Train: 100% 2000/2000 [04:31<00:00,  7.36 step/s, accuracy=0.72, loss=1.02, step=18000]  \n",
-      "Valid: 100% 5664/5667 [00:53<00:00, 105.38 uttr/s, accuracy=0.68, loss=1.51]\n",
-      "Train: 100% 2000/2000 [04:30<00:00,  7.39 step/s, accuracy=0.88, loss=0.52, step=2e+4]   \n",
-      "Valid: 100% 5664/5667 [00:54<00:00, 103.67 uttr/s, accuracy=0.70, loss=1.50]\n",
+      "Train: 100% 2000/2000 [05:18<00:00,  6.28 step/s, accuracy=1.00, loss=0.06, step=122000]  \n",
+      "Valid: 100% 5664/5667 [00:54<00:00, 103.07 uttr/s, accuracy=0.88, loss=0.68]\n",
+      "Train: 100% 2000/2000 [05:10<00:00,  6.43 step/s, accuracy=0.97, loss=0.07, step=124000] \n",
+      "Valid: 100% 5664/5667 [00:57<00:00, 98.42 uttr/s, accuracy=0.88, loss=0.68] \n",
+      "Train: 100% 2000/2000 [06:17<00:00,  5.29 step/s, accuracy=1.00, loss=0.01, step=126000] \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 105.35 uttr/s, accuracy=0.88, loss=0.67]\n",
+      "Train: 100% 2000/2000 [05:00<00:00,  6.65 step/s, accuracy=1.00, loss=0.05, step=128000]  \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 106.92 uttr/s, accuracy=0.88, loss=0.66]\n",
+      "Train: 100% 2000/2000 [04:59<00:00,  6.67 step/s, accuracy=1.00, loss=0.05, step=130000]  \n",
+      "Valid: 100% 5664/5667 [00:54<00:00, 103.12 uttr/s, accuracy=0.89, loss=0.65]\n",
       "Train:   0% 0/2000 [00:00<?, ? step/s]"
      ]
     },
@@ -526,23 +798,23 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Step 20000, best model saved. (accuracy=0.6974)\n"
+      "Step 130000, best model saved. (accuracy=0.8916)\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Train: 100% 2000/2000 [04:34<00:00,  7.29 step/s, accuracy=0.78, loss=0.67, step=22000] \n",
-      "Valid: 100% 5664/5667 [00:54<00:00, 103.27 uttr/s, accuracy=0.73, loss=1.30]\n",
-      "Train: 100% 2000/2000 [05:21<00:00,  6.22 step/s, accuracy=0.94, loss=0.38, step=24000] \n",
-      "Valid: 100% 5664/5667 [00:52<00:00, 107.76 uttr/s, accuracy=0.74, loss=1.24]\n",
-      "Train: 100% 2000/2000 [04:34<00:00,  7.29 step/s, accuracy=0.94, loss=0.32, step=26000]  \n",
-      "Valid: 100% 5664/5667 [00:54<00:00, 104.57 uttr/s, accuracy=0.73, loss=1.39]\n",
-      "Train: 100% 2000/2000 [04:34<00:00,  7.29 step/s, accuracy=0.75, loss=1.00, step=28000]  \n",
-      "Valid: 100% 5664/5667 [00:54<00:00, 103.97 uttr/s, accuracy=0.73, loss=1.30]\n",
-      "Train: 100% 2000/2000 [04:27<00:00,  7.46 step/s, accuracy=0.88, loss=0.50, step=3e+4]  \n",
-      "Valid: 100% 5664/5667 [00:54<00:00, 104.73 uttr/s, accuracy=0.74, loss=1.25]\n",
+      "Train: 100% 2000/2000 [05:22<00:00,  6.19 step/s, accuracy=1.00, loss=0.06, step=132000] \n",
+      "Valid: 100% 5664/5667 [00:55<00:00, 101.36 uttr/s, accuracy=0.88, loss=0.66]\n",
+      "Train: 100% 2000/2000 [06:04<00:00,  5.48 step/s, accuracy=0.97, loss=0.15, step=134000] \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 105.72 uttr/s, accuracy=0.88, loss=0.66]\n",
+      "Train: 100% 2000/2000 [05:19<00:00,  6.25 step/s, accuracy=0.97, loss=0.06, step=136000]  \n",
+      "Valid: 100% 5664/5667 [00:54<00:00, 104.73 uttr/s, accuracy=0.88, loss=0.65]\n",
+      "Train: 100% 2000/2000 [05:16<00:00,  6.32 step/s, accuracy=1.00, loss=0.01, step=138000]  \n",
+      "Valid: 100% 5664/5667 [00:55<00:00, 102.74 uttr/s, accuracy=0.88, loss=0.70]\n",
+      "Train: 100% 2000/2000 [28:15<00:00,  1.18 step/s, accuracy=0.94, loss=0.29, step=140000]   \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 107.28 uttr/s, accuracy=0.89, loss=0.65]\n",
       "Train:   0% 0/2000 [00:00<?, ? step/s]"
      ]
     },
@@ -550,23 +822,23 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Step 30000, best model saved. (accuracy=0.7438)\n"
+      "Step 140000, best model saved. (accuracy=0.8916)\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Train: 100% 2000/2000 [05:38<00:00,  5.91 step/s, accuracy=0.84, loss=0.55, step=32000] \n",
-      "Valid: 100% 5664/5667 [00:53<00:00, 106.59 uttr/s, accuracy=0.75, loss=1.28]\n",
-      "Train: 100% 2000/2000 [04:40<00:00,  7.12 step/s, accuracy=0.81, loss=0.52, step=34000]  \n",
-      "Valid: 100% 5664/5667 [00:54<00:00, 104.54 uttr/s, accuracy=0.77, loss=1.22]\n",
-      "Train: 100% 2000/2000 [04:36<00:00,  7.24 step/s, accuracy=0.84, loss=0.53, step=36000]  \n",
-      "Valid: 100% 5664/5667 [00:54<00:00, 104.56 uttr/s, accuracy=0.78, loss=1.12]\n",
-      "Train: 100% 2000/2000 [04:48<00:00,  6.94 step/s, accuracy=0.78, loss=0.60, step=38000] \n",
-      "Valid: 100% 5664/5667 [00:54<00:00, 103.09 uttr/s, accuracy=0.77, loss=1.18]\n",
-      "Train: 100% 2000/2000 [05:15<00:00,  6.34 step/s, accuracy=0.97, loss=0.20, step=4e+4]  \n",
-      "Valid: 100% 5664/5667 [00:52<00:00, 108.32 uttr/s, accuracy=0.78, loss=1.14]\n",
+      "Train: 100% 2000/2000 [05:14<00:00,  6.35 step/s, accuracy=1.00, loss=0.04, step=142000] \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 108.42 uttr/s, accuracy=0.89, loss=0.64]\n",
+      "Train: 100% 2000/2000 [04:28<00:00,  7.44 step/s, accuracy=1.00, loss=0.02, step=144000]  \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 107.12 uttr/s, accuracy=0.89, loss=0.66]\n",
+      "Train: 100% 2000/2000 [04:29<00:00,  7.43 step/s, accuracy=0.97, loss=0.20, step=146000] \n",
+      "Valid: 100% 5664/5667 [00:54<00:00, 103.30 uttr/s, accuracy=0.89, loss=0.67]\n",
+      "Train: 100% 2000/2000 [29:24<00:00,  1.13 step/s, accuracy=0.97, loss=0.09, step=148000]    \n",
+      "Valid: 100% 5664/5667 [03:18<00:00, 28.55 uttr/s, accuracy=0.89, loss=0.65] \n",
+      "Train: 100% 2000/2000 [05:16<00:00,  6.32 step/s, accuracy=1.00, loss=0.06, step=150000] \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 108.23 uttr/s, accuracy=0.89, loss=0.65]\n",
       "Train:   0% 0/2000 [00:00<?, ? step/s]"
      ]
     },
@@ -574,16 +846,31 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Step 40000, best model saved. (accuracy=0.7821)\n"
+      "Step 150000, best model saved. (accuracy=0.8916)\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Train: 100% 2000/2000 [04:34<00:00,  7.30 step/s, accuracy=0.78, loss=0.78, step=42000]  \n",
-      "Valid: 100% 5664/5667 [00:54<00:00, 103.96 uttr/s, accuracy=0.76, loss=1.29]\n",
-      "Train:  43% 867/2000 [01:45<02:21,  8.03 step/s, accuracy=0.91, loss=0.28, step=42866]"
+      "Train: 100% 2000/2000 [04:29<00:00,  7.41 step/s, accuracy=1.00, loss=0.01, step=152000]  \n",
+      "Valid: 100% 5664/5667 [00:52<00:00, 107.46 uttr/s, accuracy=0.88, loss=0.66]\n",
+      "Train: 100% 2000/2000 [04:30<00:00,  7.41 step/s, accuracy=1.00, loss=0.05, step=154000] \n",
+      "Valid: 100% 5664/5667 [00:53<00:00, 106.62 uttr/s, accuracy=0.89, loss=0.65]\n",
+      "Train:  14% 286/2000 [00:32<03:11,  8.97 step/s, accuracy=1.00, loss=0.02, step=154286]"
+     ]
+    },
+    {
+     "ename": "KeyboardInterrupt",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[10], line 110\u001b[0m\n\u001b[1;32m    105\u001b[0m     train_log({\n\u001b[1;32m    106\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mstatus\u001b[39m\u001b[38;5;124m'\u001b[39m: \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcompleted\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m    107\u001b[0m     })\n\u001b[1;32m    109\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;18m__name__\u001b[39m \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__main__\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m--> 110\u001b[0m     \u001b[43mmain\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparse_args\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
+      "Cell \u001b[0;32mIn[10], line 63\u001b[0m, in \u001b[0;36mmain\u001b[0;34m(data_dir, save_path, batch_size, n_workers, valid_steps, warmup_steps, total_steps, save_steps)\u001b[0m\n\u001b[1;32m     60\u001b[0m     train_iterator \u001b[38;5;241m=\u001b[39m \u001b[38;5;28miter\u001b[39m(train_loader)\n\u001b[1;32m     61\u001b[0m     batch \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mnext\u001b[39m(train_iterator)\n\u001b[0;32m---> 63\u001b[0m loss, accuracy \u001b[38;5;241m=\u001b[39m \u001b[43mmodel_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcriterion\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     64\u001b[0m batch_loss \u001b[38;5;241m=\u001b[39m loss\u001b[38;5;241m.\u001b[39mitem()\n\u001b[1;32m     65\u001b[0m batch_accuracy \u001b[38;5;241m=\u001b[39m accuracy\u001b[38;5;241m.\u001b[39mitem()\n",
+      "Cell \u001b[0;32mIn[8], line 8\u001b[0m, in \u001b[0;36mmodel_fn\u001b[0;34m(batch, model, criterion, device)\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Forward a batch through the model.\"\"\"\u001b[39;00m\n\u001b[1;32m      7\u001b[0m mels, labels \u001b[38;5;241m=\u001b[39m batch\n\u001b[0;32m----> 8\u001b[0m mels \u001b[38;5;241m=\u001b[39m \u001b[43mmels\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mto\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      9\u001b[0m labels \u001b[38;5;241m=\u001b[39m labels\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[1;32m     11\u001b[0m outs, outs_length \u001b[38;5;241m=\u001b[39m model(mels)\n",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
      ]
     }
    ],
@@ -702,7 +989,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -725,7 +1012,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 5,
    "metadata": {
     "id": "efS4pCmAJXJH"
    },
@@ -752,7 +1039,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 6,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/",
@@ -787,7 +1074,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4a12f429161a47e2832ec39d1fe0b251",
+       "model_id": "94406e7fe0a64bb187b1edb64c2c6c4a",
        "version_major": 2,
        "version_minor": 0
       },
@@ -868,7 +1155,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [
     {
@@ -878,7 +1165,7 @@
      "traceback": [
       "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
       "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[10], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mgc\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;01mdel\u001b[39;00m dataloader\n\u001b[1;32m      3\u001b[0m gc\u001b[38;5;241m.\u001b[39mcollect()\n",
+      "Cell \u001b[0;32mIn[7], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mgc\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;01mdel\u001b[39;00m dataloader\n\u001b[1;32m      3\u001b[0m gc\u001b[38;5;241m.\u001b[39mcollect()\n",
       "\u001b[0;31mNameError\u001b[0m: name 'dataloader' is not defined"
      ]
     }