{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=== Training with round  ===\n",
      "xmy22\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tiger/.local/lib/python3.11/site-packages/transformers/optimization.py:591: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ca1d3884e99843398d7a86c13af1bd38",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0, Loss: 5.0625, Accuracy: 0.0029, TrueLabel: 0.5\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 32, Loss: 1.3125, Accuracy: 0.5239, TrueLabel: 0.5\n",
      "Step 64, Loss: 0.6719, Accuracy: 0.5366, TrueLabel: 0.5\n",
      "Step 96, Loss: 0.6758, Accuracy: 0.5713, TrueLabel: 0.5\n",
      "Step 128, Loss: 0.6328, Accuracy: 0.5796, TrueLabel: 0.5\n",
      "Step 160, Loss: 0.5938, Accuracy: 0.6118, TrueLabel: 0.5\n",
      "Step 192, Loss: 0.9453, Accuracy: 0.6157, TrueLabel: 0.5\n",
      "Step 224, Loss: 0.9531, Accuracy: 0.6069, TrueLabel: 0.5\n",
      "Step 256, Loss: 0.8125, Accuracy: 0.6338, TrueLabel: 0.5\n",
      "Step 288, Loss: 0.5586, Accuracy: 0.6763, TrueLabel: 0.5\n",
      "Step 320, Loss: 0.5117, Accuracy: 0.6714, TrueLabel: 0.5\n",
      "Step 352, Loss: 0.4785, Accuracy: 0.6880, TrueLabel: 0.5\n",
      "Step 384, Loss: 0.4980, Accuracy: 0.7031, TrueLabel: 0.5\n",
      "Step 416, Loss: 0.4707, Accuracy: 0.7065, TrueLabel: 0.5\n",
      "Step 448, Loss: 0.4492, Accuracy: 0.7168, TrueLabel: 0.5\n",
      "Step 480, Loss: 0.4453, Accuracy: 0.7227, TrueLabel: 0.5\n",
      "Step 512, Loss: 0.4355, Accuracy: 0.7271, TrueLabel: 0.5\n",
      "Step 544, Loss: 0.4219, Accuracy: 0.7275, TrueLabel: 0.5\n",
      "Step 576, Loss: 0.4121, Accuracy: 0.7490, TrueLabel: 0.5\n",
      "Step 608, Loss: 0.4102, Accuracy: 0.7471, TrueLabel: 0.5\n",
      "Step 640, Loss: 0.3828, Accuracy: 0.7729, TrueLabel: 0.5\n",
      "Step 672, Loss: 0.3945, Accuracy: 0.7495, TrueLabel: 0.5\n",
      "Step 704, Loss: 0.4766, Accuracy: 0.7671, TrueLabel: 0.5\n",
      "Step 736, Loss: 0.3633, Accuracy: 0.7710, TrueLabel: 0.5\n",
      "Step 768, Loss: 0.3281, Accuracy: 0.8052, TrueLabel: 0.5\n",
      "Step 800, Loss: 0.3301, Accuracy: 0.8091, TrueLabel: 0.5\n",
      "Step 832, Loss: 0.3262, Accuracy: 0.7998, TrueLabel: 0.5\n",
      "Step 864, Loss: 0.3242, Accuracy: 0.7930, TrueLabel: 0.5\n",
      "Step 896, Loss: 0.3047, Accuracy: 0.8281, TrueLabel: 0.5\n",
      "Step 928, Loss: 0.3008, Accuracy: 0.8086, TrueLabel: 0.5\n",
      "Step 960, Loss: 0.2969, Accuracy: 0.8193, TrueLabel: 0.5\n",
      "Step 992, Loss: 0.3086, Accuracy: 0.8301, TrueLabel: 0.5\n",
      "Step 1024, Loss: 0.2988, Accuracy: 0.8335, TrueLabel: 0.5\n",
      "Step 1056, Loss: 0.2852, Accuracy: 0.8354, TrueLabel: 0.5\n",
      "Step 1088, Loss: 0.2910, Accuracy: 0.8433, TrueLabel: 0.5\n",
      "Step 1120, Loss: 0.2695, Accuracy: 0.8525, TrueLabel: 0.5\n",
      "Step 1152, Loss: 0.2832, Accuracy: 0.8433, TrueLabel: 0.5\n",
      "Step 1184, Loss: 0.2852, Accuracy: 0.8340, TrueLabel: 0.5\n",
      "Step 1216, Loss: 0.2832, Accuracy: 0.8369, TrueLabel: 0.5\n",
      "Step 1248, Loss: 0.2637, Accuracy: 0.8525, TrueLabel: 0.5\n",
      "Step 1280, Loss: 0.2461, Accuracy: 0.8604, TrueLabel: 0.5\n",
      "Step 1312, Loss: 0.2656, Accuracy: 0.8701, TrueLabel: 0.5\n",
      "Step 1344, Loss: 0.2314, Accuracy: 0.8877, TrueLabel: 0.5\n",
      "Step 1376, Loss: 0.2393, Accuracy: 0.8843, TrueLabel: 0.5\n",
      "Step 1408, Loss: 0.2305, Accuracy: 0.8784, TrueLabel: 0.5\n",
      "Step 1440, Loss: 0.2402, Accuracy: 0.8696, TrueLabel: 0.5\n",
      "Step 1472, Loss: 0.2275, Accuracy: 0.8691, TrueLabel: 0.5\n",
      "Step 1504, Loss: 0.2080, Accuracy: 0.9009, TrueLabel: 0.5\n",
      "Step 1536, Loss: 0.2070, Accuracy: 0.9087, TrueLabel: 0.5\n",
      "Step 1568, Loss: 0.2139, Accuracy: 0.8994, TrueLabel: 0.5\n",
      "Step 1600, Loss: 0.2051, Accuracy: 0.8960, TrueLabel: 0.5\n",
      "Step 1632, Loss: 0.2012, Accuracy: 0.8979, TrueLabel: 0.5\n",
      "Step 1664, Loss: 0.1904, Accuracy: 0.9097, TrueLabel: 0.5\n",
      "Step 1696, Loss: 0.2012, Accuracy: 0.9121, TrueLabel: 0.5\n",
      "Step 1728, Loss: 0.1895, Accuracy: 0.9136, TrueLabel: 0.5\n",
      "Step 1760, Loss: 0.1914, Accuracy: 0.8921, TrueLabel: 0.5\n",
      "Step 1792, Loss: 0.1924, Accuracy: 0.9102, TrueLabel: 0.5\n",
      "Step 1824, Loss: 0.1729, Accuracy: 0.9175, TrueLabel: 0.5\n",
      "Step 1856, Loss: 0.1855, Accuracy: 0.9185, TrueLabel: 0.5\n",
      "Step 1888, Loss: 0.1611, Accuracy: 0.9292, TrueLabel: 0.5\n",
      "Step 1920, Loss: 0.1895, Accuracy: 0.9111, TrueLabel: 0.5\n",
      "Step 1952, Loss: 0.1797, Accuracy: 0.9258, TrueLabel: 0.5\n",
      "Step 1984, Loss: 0.1729, Accuracy: 0.9175, TrueLabel: 0.5\n",
      "\n",
      "=== Training with relu  ===\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "99ce501de42d4e16b2bd3b7e422daaa4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0, Loss: 5.0000, Accuracy: 0.0127, TrueLabel: 0.5\n",
      "Step 32, Loss: 0.6602, Accuracy: 0.5537, TrueLabel: 0.5\n",
      "Step 64, Loss: 0.6367, Accuracy: 0.5894, TrueLabel: 0.5\n",
      "Step 96, Loss: 0.6094, Accuracy: 0.5938, TrueLabel: 0.5\n",
      "Step 128, Loss: 0.5898, Accuracy: 0.6040, TrueLabel: 0.5\n",
      "Step 160, Loss: 0.5547, Accuracy: 0.6318, TrueLabel: 0.5\n",
      "Step 192, Loss: 0.5234, Accuracy: 0.6523, TrueLabel: 0.5\n",
      "Step 224, Loss: 0.5195, Accuracy: 0.6729, TrueLabel: 0.5\n",
      "Step 256, Loss: 0.5078, Accuracy: 0.6772, TrueLabel: 0.5\n",
      "Step 288, Loss: 0.4844, Accuracy: 0.6738, TrueLabel: 0.5\n",
      "Step 320, Loss: 0.4570, Accuracy: 0.6943, TrueLabel: 0.5\n",
      "Step 352, Loss: 0.4453, Accuracy: 0.7124, TrueLabel: 0.5\n",
      "Step 384, Loss: 0.4473, Accuracy: 0.7188, TrueLabel: 0.5\n",
      "Step 416, Loss: 0.4160, Accuracy: 0.7466, TrueLabel: 0.5\n",
      "Step 448, Loss: 0.4023, Accuracy: 0.7461, TrueLabel: 0.5\n",
      "Step 480, Loss: 0.3770, Accuracy: 0.7769, TrueLabel: 0.5\n",
      "Step 512, Loss: 0.3730, Accuracy: 0.7681, TrueLabel: 0.5\n",
      "Step 544, Loss: 0.3945, Accuracy: 0.7485, TrueLabel: 0.5\n",
      "Step 576, Loss: 0.3613, Accuracy: 0.7603, TrueLabel: 0.5\n",
      "Step 608, Loss: 0.3633, Accuracy: 0.7793, TrueLabel: 0.5\n",
      "Step 640, Loss: 0.3340, Accuracy: 0.7803, TrueLabel: 0.5\n",
      "Step 672, Loss: 0.3242, Accuracy: 0.8062, TrueLabel: 0.5\n",
      "Step 704, Loss: 0.3262, Accuracy: 0.7871, TrueLabel: 0.5\n",
      "Step 736, Loss: 0.3340, Accuracy: 0.7915, TrueLabel: 0.5\n",
      "Step 768, Loss: 0.2949, Accuracy: 0.8340, TrueLabel: 0.5\n",
      "Step 800, Loss: 0.3203, Accuracy: 0.7915, TrueLabel: 0.5\n",
      "Step 832, Loss: 0.2871, Accuracy: 0.8140, TrueLabel: 0.5\n",
      "Step 864, Loss: 0.2910, Accuracy: 0.8242, TrueLabel: 0.5\n",
      "Step 896, Loss: 0.2910, Accuracy: 0.8140, TrueLabel: 0.5\n",
      "Step 928, Loss: 0.2891, Accuracy: 0.8066, TrueLabel: 0.5\n",
      "Step 960, Loss: 0.2852, Accuracy: 0.8145, TrueLabel: 0.5\n",
      "Step 992, Loss: 0.2793, Accuracy: 0.8394, TrueLabel: 0.5\n",
      "Step 1024, Loss: 0.2695, Accuracy: 0.8218, TrueLabel: 0.5\n",
      "Step 1056, Loss: 0.2988, Accuracy: 0.8125, TrueLabel: 0.5\n",
      "Step 1088, Loss: 0.2656, Accuracy: 0.8223, TrueLabel: 0.5\n",
      "Step 1120, Loss: 0.2559, Accuracy: 0.8345, TrueLabel: 0.5\n",
      "Step 1152, Loss: 0.2637, Accuracy: 0.8350, TrueLabel: 0.5\n",
      "Step 1184, Loss: 0.2598, Accuracy: 0.8359, TrueLabel: 0.5\n",
      "Step 1216, Loss: 0.2500, Accuracy: 0.8267, TrueLabel: 0.5\n",
      "Step 1248, Loss: 0.2598, Accuracy: 0.8267, TrueLabel: 0.5\n",
      "Step 1280, Loss: 0.2559, Accuracy: 0.8389, TrueLabel: 0.5\n",
      "Step 1312, Loss: 0.2598, Accuracy: 0.8379, TrueLabel: 0.5\n",
      "Step 1344, Loss: 0.2500, Accuracy: 0.8374, TrueLabel: 0.5\n",
      "Step 1376, Loss: 0.2500, Accuracy: 0.8359, TrueLabel: 0.5\n",
      "Step 1408, Loss: 0.2422, Accuracy: 0.8516, TrueLabel: 0.5\n",
      "Step 1440, Loss: 0.2471, Accuracy: 0.8394, TrueLabel: 0.5\n",
      "Step 1472, Loss: 0.2500, Accuracy: 0.8501, TrueLabel: 0.5\n",
      "Step 1504, Loss: 0.2354, Accuracy: 0.8535, TrueLabel: 0.5\n",
      "Step 1536, Loss: 0.2285, Accuracy: 0.8501, TrueLabel: 0.5\n",
      "Step 1568, Loss: 0.2383, Accuracy: 0.8511, TrueLabel: 0.5\n",
      "Step 1600, Loss: 0.2266, Accuracy: 0.8545, TrueLabel: 0.5\n",
      "Step 1632, Loss: 0.2324, Accuracy: 0.8574, TrueLabel: 0.5\n",
      "Step 1664, Loss: 0.2383, Accuracy: 0.8428, TrueLabel: 0.5\n",
      "Step 1696, Loss: 0.2363, Accuracy: 0.8389, TrueLabel: 0.5\n",
      "Step 1728, Loss: 0.2354, Accuracy: 0.8447, TrueLabel: 0.5\n",
      "Step 1760, Loss: 0.2207, Accuracy: 0.8633, TrueLabel: 0.5\n",
      "Step 1792, Loss: 0.2285, Accuracy: 0.8652, TrueLabel: 0.5\n",
      "Step 1824, Loss: 0.2178, Accuracy: 0.8579, TrueLabel: 0.5\n",
      "Step 1856, Loss: 0.2188, Accuracy: 0.8652, TrueLabel: 0.5\n",
      "Step 1888, Loss: 0.2246, Accuracy: 0.8613, TrueLabel: 0.5\n",
      "Step 1920, Loss: 0.2324, Accuracy: 0.8506, TrueLabel: 0.5\n",
      "Step 1952, Loss: 0.2295, Accuracy: 0.8638, TrueLabel: 0.5\n",
      "Step 1984, Loss: 0.2383, Accuracy: 0.8477, TrueLabel: 0.5\n",
      "\n",
      "=== Training with linear  ===\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "555dfce2a7234faa91f1f9e230b860f0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0, Loss: 5.0625, Accuracy: 0.0044, TrueLabel: 0.5\n",
      "Step 32, Loss: 1.4531, Accuracy: 0.4814, TrueLabel: 0.5\n",
      "Step 64, Loss: 0.7266, Accuracy: 0.5259, TrueLabel: 0.5\n",
      "Step 96, Loss: 0.6953, Accuracy: 0.5747, TrueLabel: 0.5\n",
      "Step 128, Loss: 0.5898, Accuracy: 0.6084, TrueLabel: 0.5\n",
      "Step 160, Loss: 0.6289, Accuracy: 0.6113, TrueLabel: 0.5\n",
      "Step 192, Loss: 0.5664, Accuracy: 0.6484, TrueLabel: 0.5\n",
      "Step 224, Loss: 0.5586, Accuracy: 0.6533, TrueLabel: 0.5\n",
      "Step 256, Loss: 0.5000, Accuracy: 0.6821, TrueLabel: 0.5\n",
      "Step 288, Loss: 0.5039, Accuracy: 0.6831, TrueLabel: 0.5\n",
      "Step 320, Loss: 0.4727, Accuracy: 0.6968, TrueLabel: 0.5\n",
      "Step 352, Loss: 0.4688, Accuracy: 0.7378, TrueLabel: 0.5\n",
      "Step 384, Loss: 0.4297, Accuracy: 0.7319, TrueLabel: 0.5\n",
      "Step 416, Loss: 0.4102, Accuracy: 0.7476, TrueLabel: 0.5\n",
      "Step 448, Loss: 0.3711, Accuracy: 0.7773, TrueLabel: 0.5\n",
      "Step 480, Loss: 0.3867, Accuracy: 0.7744, TrueLabel: 0.5\n",
      "Step 512, Loss: 0.3340, Accuracy: 0.8159, TrueLabel: 0.5\n",
      "Step 544, Loss: 0.3633, Accuracy: 0.8164, TrueLabel: 0.5\n",
      "Step 576, Loss: 0.3340, Accuracy: 0.8428, TrueLabel: 0.5\n",
      "Step 608, Loss: 0.2812, Accuracy: 0.8501, TrueLabel: 0.5\n",
      "Step 640, Loss: 0.2695, Accuracy: 0.8652, TrueLabel: 0.5\n",
      "Step 672, Loss: 0.2412, Accuracy: 0.8706, TrueLabel: 0.5\n",
      "Step 704, Loss: 0.2383, Accuracy: 0.8901, TrueLabel: 0.5\n",
      "Step 736, Loss: 0.2344, Accuracy: 0.8975, TrueLabel: 0.5\n",
      "Step 768, Loss: 0.2188, Accuracy: 0.9106, TrueLabel: 0.5\n",
      "Step 800, Loss: 0.1895, Accuracy: 0.9150, TrueLabel: 0.5\n",
      "Step 832, Loss: 0.1787, Accuracy: 0.9219, TrueLabel: 0.5\n",
      "Step 864, Loss: 0.1729, Accuracy: 0.9194, TrueLabel: 0.5\n",
      "Step 896, Loss: 0.1670, Accuracy: 0.9258, TrueLabel: 0.5\n",
      "Step 928, Loss: 0.1289, Accuracy: 0.9370, TrueLabel: 0.5\n",
      "Step 960, Loss: 0.1396, Accuracy: 0.9355, TrueLabel: 0.5\n",
      "Step 992, Loss: 0.1250, Accuracy: 0.9390, TrueLabel: 0.5\n",
      "Step 1024, Loss: 0.1348, Accuracy: 0.9473, TrueLabel: 0.5\n",
      "Step 1056, Loss: 0.1011, Accuracy: 0.9580, TrueLabel: 0.5\n",
      "Step 1088, Loss: 0.1074, Accuracy: 0.9570, TrueLabel: 0.5\n",
      "Step 1120, Loss: 0.1206, Accuracy: 0.9502, TrueLabel: 0.5\n",
      "Step 1152, Loss: 0.1216, Accuracy: 0.9487, TrueLabel: 0.5\n",
      "Step 1184, Loss: 0.1021, Accuracy: 0.9585, TrueLabel: 0.5\n",
      "Step 1216, Loss: 0.0918, Accuracy: 0.9561, TrueLabel: 0.5\n",
      "Step 1248, Loss: 0.0815, Accuracy: 0.9624, TrueLabel: 0.5\n",
      "Step 1280, Loss: 0.0854, Accuracy: 0.9683, TrueLabel: 0.5\n",
      "Step 1312, Loss: 0.0757, Accuracy: 0.9692, TrueLabel: 0.5\n",
      "Step 1344, Loss: 0.0762, Accuracy: 0.9727, TrueLabel: 0.5\n",
      "Step 1376, Loss: 0.0806, Accuracy: 0.9629, TrueLabel: 0.5\n",
      "Step 1408, Loss: 0.0459, Accuracy: 0.9819, TrueLabel: 0.5\n",
      "Step 1440, Loss: 0.0630, Accuracy: 0.9746, TrueLabel: 0.5\n",
      "Step 1472, Loss: 0.0688, Accuracy: 0.9688, TrueLabel: 0.5\n",
      "Step 1504, Loss: 0.0674, Accuracy: 0.9741, TrueLabel: 0.5\n",
      "Step 1536, Loss: 0.0791, Accuracy: 0.9722, TrueLabel: 0.5\n",
      "Step 1568, Loss: 0.0466, Accuracy: 0.9829, TrueLabel: 0.5\n",
      "Step 1600, Loss: 0.0520, Accuracy: 0.9756, TrueLabel: 0.5\n",
      "Step 1632, Loss: 0.0444, Accuracy: 0.9854, TrueLabel: 0.5\n",
      "Step 1664, Loss: 0.0471, Accuracy: 0.9819, TrueLabel: 0.5\n",
      "Step 1696, Loss: 0.0413, Accuracy: 0.9819, TrueLabel: 0.5\n",
      "Step 1728, Loss: 0.0388, Accuracy: 0.9844, TrueLabel: 0.5\n",
      "Step 1760, Loss: 0.0630, Accuracy: 0.9790, TrueLabel: 0.5\n",
      "Step 1792, Loss: 0.0359, Accuracy: 0.9893, TrueLabel: 0.5\n",
      "Step 1824, Loss: 0.0513, Accuracy: 0.9800, TrueLabel: 0.5\n",
      "Step 1856, Loss: 0.0435, Accuracy: 0.9849, TrueLabel: 0.5\n",
      "Step 1888, Loss: 0.0510, Accuracy: 0.9805, TrueLabel: 0.5\n",
      "Step 1920, Loss: 0.0288, Accuracy: 0.9907, TrueLabel: 0.5\n",
      "Step 1952, Loss: 0.0461, Accuracy: 0.9854, TrueLabel: 0.5\n",
      "Step 1984, Loss: 0.0479, Accuracy: 0.9805, TrueLabel: 0.5\n",
      "\n",
      "=== Training with exp  ===\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ed02d646c36d4f938bf64a18280db631",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0, Loss: 5.0938, Accuracy: 0.0034, TrueLabel: 0.5\n",
      "Step 32, Loss: 2.4844, Accuracy: 0.4263, TrueLabel: 0.5\n",
      "Step 64, Loss: 0.6641, Accuracy: 0.5820, TrueLabel: 0.5\n",
      "Step 96, Loss: 0.6094, Accuracy: 0.5928, TrueLabel: 0.5\n",
      "Step 128, Loss: 0.6055, Accuracy: 0.6064, TrueLabel: 0.5\n",
      "Step 160, Loss: 0.5430, Accuracy: 0.6597, TrueLabel: 0.5\n",
      "Step 192, Loss: 0.5195, Accuracy: 0.6626, TrueLabel: 0.5\n",
      "Step 224, Loss: 0.4824, Accuracy: 0.6938, TrueLabel: 0.5\n",
      "Step 256, Loss: 0.4629, Accuracy: 0.7148, TrueLabel: 0.5\n",
      "Step 288, Loss: 0.4883, Accuracy: 0.7046, TrueLabel: 0.5\n",
      "Step 320, Loss: 0.4121, Accuracy: 0.7529, TrueLabel: 0.5\n",
      "Step 352, Loss: 0.4160, Accuracy: 0.7305, TrueLabel: 0.5\n",
      "Step 384, Loss: 0.4043, Accuracy: 0.7451, TrueLabel: 0.5\n",
      "Step 416, Loss: 0.4043, Accuracy: 0.7490, TrueLabel: 0.5\n",
      "Step 448, Loss: 0.3750, Accuracy: 0.7739, TrueLabel: 0.5\n",
      "Step 480, Loss: 0.3359, Accuracy: 0.7910, TrueLabel: 0.5\n",
      "Step 512, Loss: 0.3359, Accuracy: 0.7925, TrueLabel: 0.5\n",
      "Step 544, Loss: 0.3379, Accuracy: 0.7925, TrueLabel: 0.5\n",
      "Step 576, Loss: 0.3164, Accuracy: 0.8032, TrueLabel: 0.5\n",
      "Step 608, Loss: 0.3242, Accuracy: 0.7993, TrueLabel: 0.5\n",
      "Step 640, Loss: 0.2891, Accuracy: 0.8257, TrueLabel: 0.5\n",
      "Step 672, Loss: 0.3047, Accuracy: 0.8169, TrueLabel: 0.5\n",
      "Step 704, Loss: 0.3223, Accuracy: 0.7974, TrueLabel: 0.5\n",
      "Step 736, Loss: 0.2852, Accuracy: 0.8188, TrueLabel: 0.5\n",
      "Step 768, Loss: 0.2852, Accuracy: 0.8120, TrueLabel: 0.5\n",
      "Step 800, Loss: 0.2754, Accuracy: 0.8418, TrueLabel: 0.5\n",
      "Step 832, Loss: 0.2676, Accuracy: 0.8325, TrueLabel: 0.5\n",
      "Step 864, Loss: 0.2656, Accuracy: 0.8228, TrueLabel: 0.5\n",
      "Step 896, Loss: 0.2754, Accuracy: 0.8208, TrueLabel: 0.5\n",
      "Step 928, Loss: 0.2793, Accuracy: 0.8286, TrueLabel: 0.5\n",
      "Step 960, Loss: 0.2656, Accuracy: 0.8140, TrueLabel: 0.5\n",
      "Step 992, Loss: 0.2480, Accuracy: 0.8472, TrueLabel: 0.5\n",
      "Step 1024, Loss: 0.2539, Accuracy: 0.8242, TrueLabel: 0.5\n",
      "Step 1056, Loss: 0.2471, Accuracy: 0.8491, TrueLabel: 0.5\n",
      "Step 1088, Loss: 0.2539, Accuracy: 0.8428, TrueLabel: 0.5\n",
      "Step 1120, Loss: 0.2617, Accuracy: 0.8364, TrueLabel: 0.5\n",
      "Step 1152, Loss: 0.2539, Accuracy: 0.8389, TrueLabel: 0.5\n",
      "Step 1184, Loss: 0.2275, Accuracy: 0.8521, TrueLabel: 0.5\n",
      "Step 1216, Loss: 0.2305, Accuracy: 0.8496, TrueLabel: 0.5\n",
      "Step 1248, Loss: 0.2539, Accuracy: 0.8423, TrueLabel: 0.5\n",
      "Step 1280, Loss: 0.2559, Accuracy: 0.8486, TrueLabel: 0.5\n",
      "Step 1312, Loss: 0.2363, Accuracy: 0.8472, TrueLabel: 0.5\n",
      "Step 1344, Loss: 0.2266, Accuracy: 0.8433, TrueLabel: 0.5\n",
      "Step 1376, Loss: 0.2354, Accuracy: 0.8540, TrueLabel: 0.5\n",
      "Step 1408, Loss: 0.2393, Accuracy: 0.8472, TrueLabel: 0.5\n",
      "Step 1440, Loss: 0.2285, Accuracy: 0.8501, TrueLabel: 0.5\n",
      "Step 1472, Loss: 0.2236, Accuracy: 0.8550, TrueLabel: 0.5\n",
      "Step 1504, Loss: 0.2354, Accuracy: 0.8550, TrueLabel: 0.5\n",
      "Step 1536, Loss: 0.2363, Accuracy: 0.8550, TrueLabel: 0.5\n",
      "Step 1568, Loss: 0.2041, Accuracy: 0.8721, TrueLabel: 0.5\n",
      "Step 1600, Loss: 0.2236, Accuracy: 0.8525, TrueLabel: 0.5\n",
      "Step 1632, Loss: 0.2100, Accuracy: 0.8613, TrueLabel: 0.5\n",
      "Step 1664, Loss: 0.2148, Accuracy: 0.8643, TrueLabel: 0.5\n",
      "Step 1696, Loss: 0.2324, Accuracy: 0.8428, TrueLabel: 0.5\n",
      "Step 1728, Loss: 0.2227, Accuracy: 0.8481, TrueLabel: 0.5\n",
      "Step 1760, Loss: 0.2285, Accuracy: 0.8501, TrueLabel: 0.5\n",
      "Step 1792, Loss: 0.2256, Accuracy: 0.8477, TrueLabel: 0.5\n",
      "Step 1824, Loss: 0.2207, Accuracy: 0.8535, TrueLabel: 0.5\n",
      "Step 1856, Loss: 0.2314, Accuracy: 0.8384, TrueLabel: 0.5\n",
      "Step 1888, Loss: 0.2246, Accuracy: 0.8521, TrueLabel: 0.5\n",
      "Step 1920, Loss: 0.2227, Accuracy: 0.8477, TrueLabel: 0.5\n",
      "Step 1952, Loss: 0.2139, Accuracy: 0.8594, TrueLabel: 0.5\n",
      "Step 1984, Loss: 0.2236, Accuracy: 0.8477, TrueLabel: 0.5\n",
      "\n",
      "=== Training with softmax_v2  ===\n",
      "xmy22\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "43a60d30a9fb4d96b87662bdc68bb8ea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0, Loss: 5.0000, Accuracy: 0.0015, TrueLabel: 0.5\n",
      "Step 32, Loss: 0.6289, Accuracy: 0.5723, TrueLabel: 0.5\n",
      "Step 64, Loss: 0.6172, Accuracy: 0.5830, TrueLabel: 0.5\n",
      "Step 96, Loss: 0.6016, Accuracy: 0.6104, TrueLabel: 0.5\n",
      "Step 128, Loss: 0.5312, Accuracy: 0.6519, TrueLabel: 0.5\n",
      "Step 160, Loss: 0.5117, Accuracy: 0.6729, TrueLabel: 0.5\n",
      "Step 192, Loss: 0.4824, Accuracy: 0.6958, TrueLabel: 0.5\n",
      "Step 224, Loss: 0.4570, Accuracy: 0.7129, TrueLabel: 0.5\n",
      "Step 256, Loss: 0.4551, Accuracy: 0.7183, TrueLabel: 0.5\n",
      "Step 288, Loss: 0.4453, Accuracy: 0.7036, TrueLabel: 0.5\n",
      "Step 320, Loss: 0.4043, Accuracy: 0.7256, TrueLabel: 0.5\n",
      "Step 352, Loss: 0.3945, Accuracy: 0.7490, TrueLabel: 0.5\n",
      "Step 384, Loss: 0.4258, Accuracy: 0.7266, TrueLabel: 0.5\n",
      "Step 416, Loss: 0.3965, Accuracy: 0.7446, TrueLabel: 0.5\n",
      "Step 448, Loss: 0.3555, Accuracy: 0.7705, TrueLabel: 0.5\n",
      "Step 480, Loss: 0.3691, Accuracy: 0.7734, TrueLabel: 0.5\n",
      "Step 512, Loss: 0.3477, Accuracy: 0.7842, TrueLabel: 0.5\n",
      "Step 544, Loss: 0.3633, Accuracy: 0.7739, TrueLabel: 0.5\n",
      "Step 576, Loss: 0.3438, Accuracy: 0.7881, TrueLabel: 0.5\n",
      "Step 608, Loss: 0.3496, Accuracy: 0.7793, TrueLabel: 0.5\n",
      "Step 640, Loss: 0.3379, Accuracy: 0.7949, TrueLabel: 0.5\n",
      "Step 672, Loss: 0.3477, Accuracy: 0.7744, TrueLabel: 0.5\n",
      "Step 704, Loss: 0.3516, Accuracy: 0.7954, TrueLabel: 0.5\n",
      "Step 736, Loss: 0.3379, Accuracy: 0.7832, TrueLabel: 0.5\n",
      "Step 768, Loss: 0.3359, Accuracy: 0.7949, TrueLabel: 0.5\n",
      "Step 800, Loss: 0.2988, Accuracy: 0.7935, TrueLabel: 0.5\n",
      "Step 832, Loss: 0.3340, Accuracy: 0.7837, TrueLabel: 0.5\n",
      "Step 864, Loss: 0.3184, Accuracy: 0.8008, TrueLabel: 0.5\n",
      "Step 896, Loss: 0.3223, Accuracy: 0.8071, TrueLabel: 0.5\n",
      "Step 928, Loss: 0.3242, Accuracy: 0.8071, TrueLabel: 0.5\n",
      "Step 960, Loss: 0.3262, Accuracy: 0.7954, TrueLabel: 0.5\n",
      "Step 992, Loss: 0.2949, Accuracy: 0.8252, TrueLabel: 0.5\n",
      "Step 1024, Loss: 0.3066, Accuracy: 0.8101, TrueLabel: 0.5\n",
      "Step 1056, Loss: 0.2969, Accuracy: 0.8086, TrueLabel: 0.5\n",
      "Step 1088, Loss: 0.2852, Accuracy: 0.8193, TrueLabel: 0.5\n",
      "Step 1120, Loss: 0.3184, Accuracy: 0.7881, TrueLabel: 0.5\n",
      "Step 1152, Loss: 0.3184, Accuracy: 0.7964, TrueLabel: 0.5\n",
      "Step 1184, Loss: 0.3164, Accuracy: 0.7988, TrueLabel: 0.5\n",
      "Step 1216, Loss: 0.2969, Accuracy: 0.8198, TrueLabel: 0.5\n",
      "Step 1248, Loss: 0.2676, Accuracy: 0.8301, TrueLabel: 0.5\n",
      "Step 1280, Loss: 0.2930, Accuracy: 0.8105, TrueLabel: 0.5\n",
      "Step 1312, Loss: 0.2754, Accuracy: 0.8179, TrueLabel: 0.5\n",
      "Step 1344, Loss: 0.2812, Accuracy: 0.8286, TrueLabel: 0.5\n",
      "Step 1376, Loss: 0.2832, Accuracy: 0.8232, TrueLabel: 0.5\n",
      "Step 1408, Loss: 0.2676, Accuracy: 0.8403, TrueLabel: 0.5\n",
      "Step 1440, Loss: 0.2891, Accuracy: 0.8188, TrueLabel: 0.5\n",
      "Step 1472, Loss: 0.2773, Accuracy: 0.8208, TrueLabel: 0.5\n",
      "Step 1504, Loss: 0.2852, Accuracy: 0.8223, TrueLabel: 0.5\n",
      "Step 1536, Loss: 0.2598, Accuracy: 0.8394, TrueLabel: 0.5\n",
      "Step 1568, Loss: 0.2656, Accuracy: 0.8271, TrueLabel: 0.5\n",
      "Step 1600, Loss: 0.2734, Accuracy: 0.8340, TrueLabel: 0.5\n",
      "Step 1632, Loss: 0.2715, Accuracy: 0.8315, TrueLabel: 0.5\n",
      "Step 1664, Loss: 0.2695, Accuracy: 0.8359, TrueLabel: 0.5\n",
      "Step 1696, Loss: 0.2871, Accuracy: 0.8257, TrueLabel: 0.5\n",
      "Step 1728, Loss: 0.2656, Accuracy: 0.8428, TrueLabel: 0.5\n",
      "Step 1760, Loss: 0.2520, Accuracy: 0.8433, TrueLabel: 0.5\n",
      "Step 1792, Loss: 0.2578, Accuracy: 0.8530, TrueLabel: 0.5\n",
      "Step 1824, Loss: 0.2656, Accuracy: 0.8345, TrueLabel: 0.5\n",
      "Step 1856, Loss: 0.2676, Accuracy: 0.8408, TrueLabel: 0.5\n",
      "Step 1888, Loss: 0.2422, Accuracy: 0.8511, TrueLabel: 0.5\n",
      "Step 1920, Loss: 0.2578, Accuracy: 0.8345, TrueLabel: 0.5\n",
      "Step 1952, Loss: 0.2500, Accuracy: 0.8403, TrueLabel: 0.5\n",
      "Step 1984, Loss: 0.2637, Accuracy: 0.8350, TrueLabel: 0.5\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 143024 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import AdamW, get_scheduler\n",
    "from tqdm.auto import tqdm\n",
    "import fla\n",
    "from fla.models.utils import Cache\n",
    "from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss\n",
    "from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast\n",
    "from transformers import AutoModelForCausalLM\n",
    "import random\n",
    "from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union\n",
    "from itertools import permutations\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "def create_graph(n):\n",
    "    if n % 2 != 0:\n",
    "        raise ValueError(\"n should be an even number\")\n",
    "    \n",
    "    # Step 1: Randomly divide the points into two sets S_1 and S_2\n",
    "    points = list(range(1, n + 1))\n",
    "    random.shuffle(points)\n",
    "    mid = n // 2\n",
    "    S_1, S_2 = sorted(points[:mid]), sorted(points[mid:])\n",
    "    \n",
    "    def assign_parents(S):\n",
    "        parents = {}\n",
    "        for i in range(1, len(S)):\n",
    "            possible_parents = S[:i]\n",
    "            parents[S[i]] = random.choice(possible_parents) \n",
    "        return parents\n",
    "    \n",
    "    # Step 2: Assign parent nodes within each set\n",
    "    parents_S1 = assign_parents(S_1)\n",
    "    parents_S2 = assign_parents(S_2)\n",
    "    \n",
    "    # Step 3: Build adjacency matrix\n",
    "    adjacency_matrix = np.eye(n)\n",
    "    def fill_adjacency_matrix(parents):\n",
    "        for child, parent in parents.items():\n",
    "            if parent is not None:\n",
    "                adjacency_matrix[child - 1][parent - 1] = 1\n",
    "\n",
    "    fill_adjacency_matrix(parents_S1)\n",
    "    fill_adjacency_matrix(parents_S2)\n",
    "    labels = [0 for i in range(n)]\n",
    "    if 1 in S_1:\n",
    "        for i in S_1:\n",
    "            labels[i-1] = 1\n",
    "    else:\n",
    "        for i in S_2:\n",
    "            labels[i-1] = 1\n",
    "    return labels, adjacency_matrix\n",
    "\n",
    "def generate_graph_data(num_samples=100, n=32):\n",
    "    \"\"\"\n",
    "    Generates graph data samples with reachability information.\n",
    "    \n",
    "    :param num_samples: Number of samples to generate.\n",
    "    :param n: Number of nodes in the graph.\n",
    "    :return: A list of tuples. Each tuple contains an adjacency matrix and a list of labels indicating reachability from node 1 to each node.\n",
    "    \"\"\"\n",
    "    data = []\n",
    "\n",
    "    for _ in range(num_samples):\n",
    "        labels, A = create_graph(n)\n",
    "\n",
    "        # 将邻接矩阵和label转换为张量\n",
    "        adj_matrix = torch.tensor(A, dtype=torch.float)\n",
    "        # adj_matrix  =  adj_matrix.transpose(0,1)\n",
    "\n",
    "        labels = torch.tensor(labels, dtype=torch.long)\n",
    "\n",
    "        # 添加到数据中\n",
    "        data.append((adj_matrix, labels))\n",
    "    return data\n",
    "\n",
    "# 自定义 Dataset 类\n",
    "class S5Dataset(Dataset):\n",
    "    def __init__(self, k=4, num_samples=3):\n",
    "        self.k = k\n",
    "        self.num_samples = num_samples\n",
    "        self.data = generate_graph_data( self.num_samples,self.k)\n",
    "    def __len__(self):\n",
    "        return self.num_samples  # 每次只返回 num_samples 条数据\n",
    "    def data_reset(self):\n",
    "        self.data = generate_graph_data( self.num_samples,self.k)\n",
    "    def __getitem__(self, idx):\n",
    "        # 每次访问时重新生成数据，保证每个 epoch 数据不同\n",
    "        \n",
    "        input_ids, labels = self.data[idx]\n",
    "        return {\n",
    "            \"input_ids\": input_ids,\n",
    "            \"labels\": labels\n",
    "        }\n",
    "\n",
    "\n",
    "dtype = torch.bfloat16\n",
    "device = \"cuda:1\" if torch.cuda.is_available() else \"cpu\"\n",
    "class Emb(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.hidden_size = config.hidden_size\n",
    "    def forward(self, x):\n",
    "        # x shape: (batch_size, seq_len, input_dim)\n",
    "        batch_size, seq_len, input_dim = x.shape\n",
    "\n",
    "        # Step 1: 构建 one-hot 形式的位置编码（只在第一个位置加）\n",
    "        # 创建一个长度为 seq_len 的 one-hot 向量，只有第一个位置是 1\n",
    "        pos_onehot = torch.zeros(seq_len, seq_len, device=x.device)\n",
    "        pos_onehot[0, 0] = 1  # 第一个位置为 one-hot\n",
    "        pos_emb = pos_onehot.unsqueeze(0).expand(batch_size, -1, -1)  # (batch_size, seq_len, seq_len)\n",
    "\n",
    "        # # Step 2: 拼接到输入 x 的最后一个维度\n",
    "        # x = torch.cat([x, pos_emb], dim=-1)  # (batch_size, seq_len, input_dim + seq_len)\n",
    "\n",
    "        # Step 3: 如果拼接后仍小于 hidden_size，则继续用 zero padding 补齐\n",
    "        current_dim = x.size(-1)\n",
    "        if current_dim < self.hidden_size:\n",
    "            pad_size = list(x.shape)\n",
    "            pad_size[-1] = self.hidden_size - current_dim\n",
    "            padding = torch.zeros(*pad_size, device=x.device)\n",
    "            x = torch.cat([x, padding], dim=-1)  # (batch_size, seq_len, hidden_size)\n",
    "        else:\n",
    "            assert False\n",
    "        return x.to(dtype) \n",
    "        return torch.nn.functional.normalize(x+1e-5,dim=-1).to(dtype)\n",
    "    \n",
    "class Iden(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# 训练函数\n",
    "def train_model_with_layers(num_layers, run_name=\"run\", log_dict=None,length = 64,kernel = 'round'):\n",
    "# 参数设置\n",
    "    kl2 = True\n",
    "    # if kernel == 'linear':\n",
    "    #     kl2 = True\n",
    "    config = fla.models.DeltaTransformerConfig(\n",
    "            vocab_size=128,\n",
    "            hidden_size=1024,\n",
    "            num_hidden_layers=num_layers,\n",
    "            num_heads=8,\n",
    "            num_kv_heads=8,\n",
    "            max_position_embeddings=2048,\n",
    "            layer_norm_eps=1e-6,\n",
    "            # use_rotary_pos_emb=None,\n",
    "            fuse_cross_entropy = False,\n",
    "            rope_theta = -1,\n",
    "            kl2 = kl2,\n",
    "            kernel = kernel,\n",
    "            usequ = True,\n",
    "            bos_token_id = 127,\n",
    "            eos_token_id = 126,\n",
    "            returnu = False, \n",
    "            fuse_norm= False,\n",
    "            fuse_swiglu = False,\n",
    "        )\n",
    "\n",
    "    # config = fla.models.TransformerConfig(\n",
    "    #     vocab_size=128,\n",
    "    #     hidden_size=1024,\n",
    "    #     num_hidden_layers=num_layers,\n",
    "    #     num_heads=8,\n",
    "    #     num_kv_heads=8,\n",
    "    #     max_position_embeddings=2048,\n",
    "    #     layer_norm_eps=1e-6,\n",
    "    #     use_rotary_pos_emb=None,\n",
    "    #     fuse_cross_entropy = False,\n",
    "    #     rope_theta = 1e4,\n",
    "    # )\n",
    "    model = AutoModelForCausalLM.from_config(config)\n",
    "    model.model.embeddings = Emb(config)    \n",
    "    # model.model.layers[0].mlp = Iden(config)\n",
    "    model.to(device).to(dtype)\n",
    "\n",
    "    # 超参数\n",
    "    num_epochs = 2000\n",
    "    batch_size = 32\n",
    "    k = int(length)\n",
    "    num_samples_per_batch = batch_size  # 每次生成的样本数 = batch size\n",
    "\n",
    "    # 数据集和数据加载器\n",
    "    dataset = S5Dataset(k=k, num_samples=num_samples_per_batch)\n",
    "    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    # 优化器和学习率调度器\n",
    "    optimizer = AdamW(model.parameters(), lr=1e-3,weight_decay=0.1)\n",
    "    num_training_steps = num_epochs * len(dataloader)\n",
    "    lr_scheduler = get_scheduler(\n",
    "        name=\"linear\", optimizer=optimizer, num_warmup_steps=128, num_training_steps=num_training_steps\n",
    "    )\n",
    "\n",
    "    # 训练循环\n",
    "    progress_bar = tqdm(range(num_training_steps))\n",
    "\n",
    "    model.train()\n",
    "    step = 0\n",
    "    acc_log = []\n",
    "    loss_log = []\n",
    "    for epoch in range(num_epochs):\n",
    "        dataloader.dataset.data_reset()\n",
    "        for batch in dataloader:\n",
    "            input_ids = batch[\"input_ids\"].to(device)\n",
    "            labels = batch[\"labels\"].to(device)\n",
    "            outputs = model(\n",
    "                input_ids=input_ids,\n",
    "                labels=labels,\n",
    "                return_dict = True\n",
    "            )\n",
    "\n",
    "            loss = outputs.loss\n",
    "            # 准确率计算\n",
    "            logits = outputs.logits\n",
    "            # print(outputs)\n",
    "            # print(input_ids[0])\n",
    "            # print(labels[0])\n",
    "            # assert 1==2\n",
    "            predictions = torch.argmax(logits, dim=-1)\n",
    "            correct = (predictions == labels).sum().item()\n",
    "\n",
    "            \n",
    "            total = labels.numel()\n",
    "            accuracy = correct / total\n",
    "\n",
    "            # 反向传播\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # 打印信息\n",
    "            if step%32==0:\n",
    "                num = labels.float().mean().item()\n",
    "                print(f\"Step {step}, Loss: {loss.item():.4f}, Accuracy: {accuracy:.4f}, TrueLabel: {num}\")\n",
    "            step += 1\n",
    "\n",
    "            progress_bar.update(1)\n",
    "            acc_log.append(accuracy)\n",
    "            loss_log.append(loss.item())\n",
    "        # if accuracy>0.99:\n",
    "        #     k*=2\n",
    "        #     print(f'setp {step}, length come to {k}')\n",
    "        #     dataset = S5Dataset(k=k, num_samples=num_samples_per_batch)\n",
    "        #     dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "    log_dict[run_name] = {\"acc\": acc_log, \"loss\": loss_log}\n",
    "if __name__ == \"__main__\":\n",
    "    log_dict = {}\n",
    "\n",
    "    # 不同层数跑实验\n",
    "    for kernel in ['round','relu','linear','exp','softmax_v2']:\n",
    "        print(f\"\\n=== Training with {kernel}  ===\")\n",
    "        train_model_with_layers(1, run_name=f\"{kernel}\", log_dict=log_dict,kernel = kernel)\n",
    "    # train_model_with_layers(layers, run_name=f\"{round}\", log_dict=log_dict,kernel = 'round')\n",
    "    # 绘图\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    for key in log_dict:\n",
    "        accs = log_dict[key]['acc']\n",
    "        steps = np.arange(len(accs))\n",
    "        if 'softmax' in key:\n",
    "            key = 'softmax'\n",
    "        plt.plot(steps, accs, label=key)\n",
    "\n",
    "    # plt.title(\"Accuracy vs Training Steps for Different Layer Counts\")\n",
    "    plt.xlabel(\"Training Step\")\n",
    "    plt.ylabel(\"Accuracy\")\n",
    "\n",
    "    plt.ylim(0.5,1)\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.savefig(\"accuracy_comparison.pdf\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=== Training with 1 layers ===\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tiger/.local/lib/python3.11/site-packages/transformers/optimization.py:591: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9a5eb55439e544cfacd6fd24725b4a38",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0, Loss: 5.2500, Accuracy: 0.0054, TrueLabel: 0.5\n",
      "Step 32, Loss: 0.6445, Accuracy: 0.5801, TrueLabel: 0.5\n",
      "Step 64, Loss: 0.5977, Accuracy: 0.6035, TrueLabel: 0.5\n",
      "Step 96, Loss: 0.5547, Accuracy: 0.6436, TrueLabel: 0.5\n",
      "Step 128, Loss: 0.5195, Accuracy: 0.6685, TrueLabel: 0.5\n",
      "Step 160, Loss: 0.4844, Accuracy: 0.6743, TrueLabel: 0.5\n",
      "Step 192, Loss: 0.4434, Accuracy: 0.7197, TrueLabel: 0.5\n",
      "Step 224, Loss: 0.4453, Accuracy: 0.7046, TrueLabel: 0.5\n",
      "Step 256, Loss: 0.4238, Accuracy: 0.7300, TrueLabel: 0.5\n",
      "Step 288, Loss: 0.4180, Accuracy: 0.7427, TrueLabel: 0.5\n",
      "Step 320, Loss: 0.4199, Accuracy: 0.7368, TrueLabel: 0.5\n",
      "Step 352, Loss: 0.4023, Accuracy: 0.7373, TrueLabel: 0.5\n",
      "Step 384, Loss: 0.4043, Accuracy: 0.7407, TrueLabel: 0.5\n",
      "Step 416, Loss: 0.3984, Accuracy: 0.7451, TrueLabel: 0.5\n",
      "Step 448, Loss: 0.3887, Accuracy: 0.7500, TrueLabel: 0.5\n",
      "Step 480, Loss: 0.3828, Accuracy: 0.7588, TrueLabel: 0.5\n",
      "Step 512, Loss: 0.3750, Accuracy: 0.7568, TrueLabel: 0.5\n",
      "Step 544, Loss: 0.3691, Accuracy: 0.7588, TrueLabel: 0.5\n",
      "Step 576, Loss: 0.3574, Accuracy: 0.7705, TrueLabel: 0.5\n",
      "Step 608, Loss: 0.3789, Accuracy: 0.7642, TrueLabel: 0.5\n",
      "Step 640, Loss: 0.3750, Accuracy: 0.7461, TrueLabel: 0.5\n",
      "Step 672, Loss: 0.3574, Accuracy: 0.7778, TrueLabel: 0.5\n",
      "Step 704, Loss: 0.3457, Accuracy: 0.7764, TrueLabel: 0.5\n",
      "Step 736, Loss: 0.3574, Accuracy: 0.7759, TrueLabel: 0.5\n",
      "Step 768, Loss: 0.3750, Accuracy: 0.7651, TrueLabel: 0.5\n",
      "Step 800, Loss: 0.3594, Accuracy: 0.7563, TrueLabel: 0.5\n",
      "Step 832, Loss: 0.3496, Accuracy: 0.7666, TrueLabel: 0.5\n",
      "Step 864, Loss: 0.3242, Accuracy: 0.7832, TrueLabel: 0.5\n",
      "Step 896, Loss: 0.3438, Accuracy: 0.7783, TrueLabel: 0.5\n",
      "Step 928, Loss: 0.3535, Accuracy: 0.7788, TrueLabel: 0.5\n",
      "Step 960, Loss: 0.3516, Accuracy: 0.7783, TrueLabel: 0.5\n",
      "Step 992, Loss: 0.3457, Accuracy: 0.7778, TrueLabel: 0.5\n",
      "Step 1024, Loss: 0.3359, Accuracy: 0.7720, TrueLabel: 0.5\n",
      "Step 1056, Loss: 0.3457, Accuracy: 0.7686, TrueLabel: 0.5\n",
      "Step 1088, Loss: 0.3359, Accuracy: 0.7905, TrueLabel: 0.5\n",
      "Step 1120, Loss: 0.3398, Accuracy: 0.7949, TrueLabel: 0.5\n",
      "Step 1152, Loss: 0.3320, Accuracy: 0.7959, TrueLabel: 0.5\n",
      "Step 1184, Loss: 0.3301, Accuracy: 0.7822, TrueLabel: 0.5\n",
      "Step 1216, Loss: 0.3359, Accuracy: 0.7871, TrueLabel: 0.5\n",
      "Step 1248, Loss: 0.3555, Accuracy: 0.7793, TrueLabel: 0.5\n",
      "Step 1280, Loss: 0.3398, Accuracy: 0.7847, TrueLabel: 0.5\n",
      "Step 1312, Loss: 0.3379, Accuracy: 0.7710, TrueLabel: 0.5\n",
      "Step 1344, Loss: 0.3398, Accuracy: 0.7847, TrueLabel: 0.5\n",
      "Step 1376, Loss: 0.3438, Accuracy: 0.7651, TrueLabel: 0.5\n",
      "Step 1408, Loss: 0.3320, Accuracy: 0.7871, TrueLabel: 0.5\n",
      "Step 1440, Loss: 0.3457, Accuracy: 0.7783, TrueLabel: 0.5\n",
      "Step 1472, Loss: 0.3359, Accuracy: 0.7788, TrueLabel: 0.5\n",
      "Step 1504, Loss: 0.3242, Accuracy: 0.8042, TrueLabel: 0.5\n",
      "Step 1536, Loss: 0.3223, Accuracy: 0.8125, TrueLabel: 0.5\n",
      "Step 1568, Loss: 0.3203, Accuracy: 0.7891, TrueLabel: 0.5\n",
      "Step 1600, Loss: 0.3516, Accuracy: 0.7739, TrueLabel: 0.5\n",
      "Step 1632, Loss: 0.3359, Accuracy: 0.7847, TrueLabel: 0.5\n",
      "Step 1664, Loss: 0.3398, Accuracy: 0.7891, TrueLabel: 0.5\n",
      "Step 1696, Loss: 0.3242, Accuracy: 0.7900, TrueLabel: 0.5\n",
      "Step 1728, Loss: 0.3594, Accuracy: 0.7661, TrueLabel: 0.5\n",
      "Step 1760, Loss: 0.3418, Accuracy: 0.7827, TrueLabel: 0.5\n",
      "Step 1792, Loss: 0.3574, Accuracy: 0.7612, TrueLabel: 0.5\n",
      "Step 1824, Loss: 0.3398, Accuracy: 0.7769, TrueLabel: 0.5\n",
      "Step 1856, Loss: 0.3359, Accuracy: 0.7979, TrueLabel: 0.5\n",
      "Step 1888, Loss: 0.3398, Accuracy: 0.7886, TrueLabel: 0.5\n",
      "Step 1920, Loss: 0.3359, Accuracy: 0.7827, TrueLabel: 0.5\n",
      "Step 1952, Loss: 0.3262, Accuracy: 0.7935, TrueLabel: 0.5\n",
      "Step 1984, Loss: 0.3398, Accuracy: 0.7949, TrueLabel: 0.5\n",
      "Step 2016, Loss: 0.3125, Accuracy: 0.8076, TrueLabel: 0.5\n",
      "Step 2048, Loss: 0.3418, Accuracy: 0.7876, TrueLabel: 0.5\n",
      "Step 2080, Loss: 0.3203, Accuracy: 0.7964, TrueLabel: 0.5\n",
      "Step 2112, Loss: 0.3242, Accuracy: 0.7979, TrueLabel: 0.5\n",
      "Step 2144, Loss: 0.3242, Accuracy: 0.7896, TrueLabel: 0.5\n",
      "Step 2176, Loss: 0.3184, Accuracy: 0.7969, TrueLabel: 0.5\n",
      "Step 2208, Loss: 0.3242, Accuracy: 0.7974, TrueLabel: 0.5\n",
      "Step 2240, Loss: 0.3203, Accuracy: 0.7944, TrueLabel: 0.5\n",
      "Step 2272, Loss: 0.3066, Accuracy: 0.7969, TrueLabel: 0.5\n",
      "Step 2304, Loss: 0.3184, Accuracy: 0.8022, TrueLabel: 0.5\n",
      "Step 2336, Loss: 0.3203, Accuracy: 0.7881, TrueLabel: 0.5\n",
      "Step 2368, Loss: 0.3262, Accuracy: 0.7969, TrueLabel: 0.5\n",
      "Step 2400, Loss: 0.3340, Accuracy: 0.7832, TrueLabel: 0.5\n",
      "Step 2432, Loss: 0.3223, Accuracy: 0.7896, TrueLabel: 0.5\n",
      "Step 2464, Loss: 0.3066, Accuracy: 0.7935, TrueLabel: 0.5\n",
      "Step 2496, Loss: 0.3301, Accuracy: 0.7886, TrueLabel: 0.5\n",
      "Step 2528, Loss: 0.3008, Accuracy: 0.8101, TrueLabel: 0.5\n",
      "Step 2560, Loss: 0.2988, Accuracy: 0.7983, TrueLabel: 0.5\n",
      "Step 2592, Loss: 0.3105, Accuracy: 0.7988, TrueLabel: 0.5\n",
      "Step 2624, Loss: 0.3125, Accuracy: 0.8071, TrueLabel: 0.5\n",
      "Step 2656, Loss: 0.3203, Accuracy: 0.7798, TrueLabel: 0.5\n",
      "Step 2688, Loss: 0.3086, Accuracy: 0.8037, TrueLabel: 0.5\n",
      "Step 2720, Loss: 0.3105, Accuracy: 0.8086, TrueLabel: 0.5\n",
      "Step 2752, Loss: 0.3105, Accuracy: 0.8047, TrueLabel: 0.5\n",
      "Step 2784, Loss: 0.3066, Accuracy: 0.8076, TrueLabel: 0.5\n",
      "Step 2816, Loss: 0.3203, Accuracy: 0.7979, TrueLabel: 0.5\n",
      "Step 2848, Loss: 0.3145, Accuracy: 0.7993, TrueLabel: 0.5\n",
      "Step 2880, Loss: 0.3027, Accuracy: 0.8047, TrueLabel: 0.5\n",
      "Step 2912, Loss: 0.3105, Accuracy: 0.7998, TrueLabel: 0.5\n",
      "Step 2944, Loss: 0.3008, Accuracy: 0.7993, TrueLabel: 0.5\n",
      "Step 2976, Loss: 0.3223, Accuracy: 0.7993, TrueLabel: 0.5\n",
      "Step 3008, Loss: 0.3242, Accuracy: 0.7881, TrueLabel: 0.5\n",
      "Step 3040, Loss: 0.3145, Accuracy: 0.8076, TrueLabel: 0.5\n",
      "Step 3072, Loss: 0.3066, Accuracy: 0.7905, TrueLabel: 0.5\n",
      "Step 3104, Loss: 0.3066, Accuracy: 0.8115, TrueLabel: 0.5\n",
      "Step 3136, Loss: 0.3086, Accuracy: 0.8008, TrueLabel: 0.5\n",
      "Step 3168, Loss: 0.3145, Accuracy: 0.8076, TrueLabel: 0.5\n",
      "Step 3200, Loss: 0.3223, Accuracy: 0.7842, TrueLabel: 0.5\n",
      "Step 3232, Loss: 0.3184, Accuracy: 0.7915, TrueLabel: 0.5\n",
      "Step 3264, Loss: 0.3066, Accuracy: 0.8042, TrueLabel: 0.5\n",
      "Step 3296, Loss: 0.3066, Accuracy: 0.8066, TrueLabel: 0.5\n",
      "Step 3328, Loss: 0.3184, Accuracy: 0.8052, TrueLabel: 0.5\n",
      "Step 3360, Loss: 0.3027, Accuracy: 0.8062, TrueLabel: 0.5\n",
      "Step 3392, Loss: 0.3164, Accuracy: 0.8120, TrueLabel: 0.5\n",
      "Step 3424, Loss: 0.3125, Accuracy: 0.8052, TrueLabel: 0.5\n",
      "Step 3456, Loss: 0.3203, Accuracy: 0.7852, TrueLabel: 0.5\n",
      "Step 3488, Loss: 0.3008, Accuracy: 0.7979, TrueLabel: 0.5\n",
      "Step 3520, Loss: 0.3301, Accuracy: 0.7964, TrueLabel: 0.5\n",
      "Step 3552, Loss: 0.3066, Accuracy: 0.8052, TrueLabel: 0.5\n",
      "Step 3584, Loss: 0.3086, Accuracy: 0.7959, TrueLabel: 0.5\n",
      "Step 3616, Loss: 0.2988, Accuracy: 0.8130, TrueLabel: 0.5\n",
      "Step 3648, Loss: 0.3105, Accuracy: 0.8120, TrueLabel: 0.5\n",
      "Step 3680, Loss: 0.3125, Accuracy: 0.8057, TrueLabel: 0.5\n",
      "Step 3712, Loss: 0.3105, Accuracy: 0.7993, TrueLabel: 0.5\n",
      "Step 3744, Loss: 0.3125, Accuracy: 0.8091, TrueLabel: 0.5\n",
      "Step 3776, Loss: 0.3320, Accuracy: 0.7979, TrueLabel: 0.5\n",
      "Step 3808, Loss: 0.3086, Accuracy: 0.8076, TrueLabel: 0.5\n",
      "Step 3840, Loss: 0.3184, Accuracy: 0.8008, TrueLabel: 0.5\n",
      "Step 3872, Loss: 0.3008, Accuracy: 0.8276, TrueLabel: 0.5\n",
      "Step 3904, Loss: 0.3066, Accuracy: 0.8047, TrueLabel: 0.5\n",
      "Step 3936, Loss: 0.2969, Accuracy: 0.8105, TrueLabel: 0.5\n",
      "Step 3968, Loss: 0.3184, Accuracy: 0.7993, TrueLabel: 0.5\n",
      "Step 4000, Loss: 0.3047, Accuracy: 0.8057, TrueLabel: 0.5\n",
      "Step 4032, Loss: 0.3027, Accuracy: 0.8135, TrueLabel: 0.5\n",
      "Step 4064, Loss: 0.3223, Accuracy: 0.7983, TrueLabel: 0.5\n",
      "Step 4096, Loss: 0.3086, Accuracy: 0.8115, TrueLabel: 0.5\n",
      "Step 4128, Loss: 0.3066, Accuracy: 0.7915, TrueLabel: 0.5\n",
      "Step 4160, Loss: 0.3145, Accuracy: 0.8013, TrueLabel: 0.5\n",
      "Step 4192, Loss: 0.3008, Accuracy: 0.8120, TrueLabel: 0.5\n",
      "Step 4224, Loss: 0.3027, Accuracy: 0.8164, TrueLabel: 0.5\n",
      "Step 4256, Loss: 0.3145, Accuracy: 0.8032, TrueLabel: 0.5\n",
      "Step 4288, Loss: 0.3164, Accuracy: 0.7983, TrueLabel: 0.5\n",
      "Step 4320, Loss: 0.2988, Accuracy: 0.8154, TrueLabel: 0.5\n",
      "Step 4352, Loss: 0.3203, Accuracy: 0.8018, TrueLabel: 0.5\n",
      "Step 4384, Loss: 0.3086, Accuracy: 0.8057, TrueLabel: 0.5\n",
      "Step 4416, Loss: 0.3164, Accuracy: 0.7935, TrueLabel: 0.5\n",
      "Step 4448, Loss: 0.3027, Accuracy: 0.8115, TrueLabel: 0.5\n",
      "Step 4480, Loss: 0.3086, Accuracy: 0.8076, TrueLabel: 0.5\n",
      "Step 4512, Loss: 0.3027, Accuracy: 0.8096, TrueLabel: 0.5\n",
      "Step 4544, Loss: 0.3105, Accuracy: 0.7979, TrueLabel: 0.5\n",
      "Step 4576, Loss: 0.3047, Accuracy: 0.8042, TrueLabel: 0.5\n",
      "Step 4608, Loss: 0.3105, Accuracy: 0.8047, TrueLabel: 0.5\n",
      "Step 4640, Loss: 0.3008, Accuracy: 0.8096, TrueLabel: 0.5\n",
      "Step 4672, Loss: 0.3184, Accuracy: 0.8013, TrueLabel: 0.5\n",
      "Step 4704, Loss: 0.3281, Accuracy: 0.7964, TrueLabel: 0.5\n",
      "Step 4736, Loss: 0.3262, Accuracy: 0.7837, TrueLabel: 0.5\n",
      "Step 4768, Loss: 0.3047, Accuracy: 0.8062, TrueLabel: 0.5\n",
      "Step 4800, Loss: 0.3086, Accuracy: 0.7915, TrueLabel: 0.5\n",
      "Step 4832, Loss: 0.2988, Accuracy: 0.8135, TrueLabel: 0.5\n",
      "Step 4864, Loss: 0.3047, Accuracy: 0.8101, TrueLabel: 0.5\n",
      "Step 4896, Loss: 0.3242, Accuracy: 0.7925, TrueLabel: 0.5\n",
      "Step 4928, Loss: 0.3301, Accuracy: 0.7881, TrueLabel: 0.5\n",
      "Step 4960, Loss: 0.3066, Accuracy: 0.8071, TrueLabel: 0.5\n",
      "Step 4992, Loss: 0.3105, Accuracy: 0.7988, TrueLabel: 0.5\n",
      "\n",
      "=== Training with 2 layers ===\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tiger/.local/lib/python3.11/site-packages/transformers/optimization.py:591: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7a8c5baae29c4c65a4a69f6cd71ae278",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0, Loss: 4.5625, Accuracy: 0.0039, TrueLabel: 0.5\n",
      "Step 32, Loss: 0.6484, Accuracy: 0.5454, TrueLabel: 0.5\n",
      "Step 64, Loss: 0.6367, Accuracy: 0.5698, TrueLabel: 0.5\n",
      "Step 96, Loss: 0.5430, Accuracy: 0.6445, TrueLabel: 0.5\n",
      "Step 128, Loss: 0.5117, Accuracy: 0.6797, TrueLabel: 0.5\n",
      "Step 160, Loss: 0.4727, Accuracy: 0.6851, TrueLabel: 0.5\n",
      "Step 192, Loss: 0.4688, Accuracy: 0.6851, TrueLabel: 0.5\n",
      "Step 224, Loss: 0.4512, Accuracy: 0.7183, TrueLabel: 0.5\n",
      "Step 256, Loss: 0.4492, Accuracy: 0.7280, TrueLabel: 0.5\n",
      "Step 288, Loss: 0.4473, Accuracy: 0.7148, TrueLabel: 0.5\n",
      "Step 320, Loss: 0.4492, Accuracy: 0.7095, TrueLabel: 0.5\n",
      "Step 352, Loss: 0.3945, Accuracy: 0.7246, TrueLabel: 0.5\n",
      "Step 384, Loss: 0.3887, Accuracy: 0.7642, TrueLabel: 0.5\n",
      "Step 416, Loss: 0.3965, Accuracy: 0.7456, TrueLabel: 0.5\n",
      "Step 448, Loss: 0.4023, Accuracy: 0.7422, TrueLabel: 0.5\n",
      "Step 480, Loss: 0.4453, Accuracy: 0.7354, TrueLabel: 0.5\n",
      "Step 512, Loss: 0.3789, Accuracy: 0.7646, TrueLabel: 0.5\n",
      "Step 544, Loss: 0.3711, Accuracy: 0.7661, TrueLabel: 0.5\n",
      "Step 576, Loss: 0.3672, Accuracy: 0.7642, TrueLabel: 0.5\n",
      "Step 608, Loss: 0.3945, Accuracy: 0.7407, TrueLabel: 0.5\n",
      "Step 640, Loss: 0.3906, Accuracy: 0.7612, TrueLabel: 0.5\n",
      "Step 672, Loss: 0.3711, Accuracy: 0.7583, TrueLabel: 0.5\n",
      "Step 704, Loss: 0.3633, Accuracy: 0.7720, TrueLabel: 0.5\n",
      "Step 736, Loss: 0.3809, Accuracy: 0.7568, TrueLabel: 0.5\n",
      "Step 768, Loss: 0.3652, Accuracy: 0.7681, TrueLabel: 0.5\n",
      "Step 800, Loss: 0.3613, Accuracy: 0.7856, TrueLabel: 0.5\n",
      "Step 832, Loss: 0.3477, Accuracy: 0.7881, TrueLabel: 0.5\n",
      "Step 864, Loss: 0.3535, Accuracy: 0.7808, TrueLabel: 0.5\n",
      "Step 896, Loss: 0.3594, Accuracy: 0.7705, TrueLabel: 0.5\n",
      "Step 928, Loss: 0.3789, Accuracy: 0.7490, TrueLabel: 0.5\n",
      "Step 960, Loss: 0.3574, Accuracy: 0.7529, TrueLabel: 0.5\n",
      "Step 992, Loss: 0.3496, Accuracy: 0.7778, TrueLabel: 0.5\n",
      "Step 1024, Loss: 0.3613, Accuracy: 0.7627, TrueLabel: 0.5\n",
      "Step 1056, Loss: 0.3672, Accuracy: 0.7534, TrueLabel: 0.5\n",
      "Step 1088, Loss: 0.3672, Accuracy: 0.7676, TrueLabel: 0.5\n",
      "Step 1120, Loss: 0.3594, Accuracy: 0.7651, TrueLabel: 0.5\n",
      "Step 1152, Loss: 0.3633, Accuracy: 0.7734, TrueLabel: 0.5\n",
      "Step 1184, Loss: 0.3457, Accuracy: 0.7817, TrueLabel: 0.5\n",
      "Step 1216, Loss: 0.3496, Accuracy: 0.7681, TrueLabel: 0.5\n",
      "Step 1248, Loss: 0.3594, Accuracy: 0.7573, TrueLabel: 0.5\n",
      "Step 1280, Loss: 0.3301, Accuracy: 0.7866, TrueLabel: 0.5\n",
      "Step 1312, Loss: 0.3770, Accuracy: 0.7534, TrueLabel: 0.5\n",
      "Step 1344, Loss: 0.3457, Accuracy: 0.7866, TrueLabel: 0.5\n",
      "Step 1376, Loss: 0.3672, Accuracy: 0.7656, TrueLabel: 0.5\n",
      "Step 1408, Loss: 0.3574, Accuracy: 0.7593, TrueLabel: 0.5\n",
      "Step 1440, Loss: 0.3359, Accuracy: 0.7910, TrueLabel: 0.5\n",
      "Step 1472, Loss: 0.3281, Accuracy: 0.8022, TrueLabel: 0.5\n",
      "Step 1504, Loss: 0.3438, Accuracy: 0.7876, TrueLabel: 0.5\n",
      "Step 1536, Loss: 0.3457, Accuracy: 0.7778, TrueLabel: 0.5\n",
      "Step 1568, Loss: 0.3477, Accuracy: 0.7856, TrueLabel: 0.5\n",
      "Step 1600, Loss: 0.3516, Accuracy: 0.7720, TrueLabel: 0.5\n",
      "Step 1632, Loss: 0.3535, Accuracy: 0.7734, TrueLabel: 0.5\n",
      "Step 1664, Loss: 0.3770, Accuracy: 0.7510, TrueLabel: 0.5\n",
      "Step 1696, Loss: 0.3594, Accuracy: 0.7695, TrueLabel: 0.5\n",
      "Step 1728, Loss: 0.3477, Accuracy: 0.7739, TrueLabel: 0.5\n",
      "Step 1760, Loss: 0.3711, Accuracy: 0.7607, TrueLabel: 0.5\n",
      "Step 1792, Loss: 0.3574, Accuracy: 0.7651, TrueLabel: 0.5\n",
      "Step 1824, Loss: 0.3320, Accuracy: 0.7974, TrueLabel: 0.5\n",
      "Step 1856, Loss: 0.3340, Accuracy: 0.7949, TrueLabel: 0.5\n",
      "Step 1888, Loss: 0.3359, Accuracy: 0.7886, TrueLabel: 0.5\n",
      "Step 1920, Loss: 0.3340, Accuracy: 0.7866, TrueLabel: 0.5\n",
      "Step 1952, Loss: 0.3457, Accuracy: 0.7783, TrueLabel: 0.5\n",
      "Step 1984, Loss: 0.3477, Accuracy: 0.7764, TrueLabel: 0.5\n",
      "Step 2016, Loss: 0.3398, Accuracy: 0.7886, TrueLabel: 0.5\n",
      "Step 2048, Loss: 0.3379, Accuracy: 0.7852, TrueLabel: 0.5\n",
      "Step 2080, Loss: 0.3457, Accuracy: 0.7822, TrueLabel: 0.5\n",
      "Step 2112, Loss: 0.3496, Accuracy: 0.7808, TrueLabel: 0.5\n",
      "Step 2144, Loss: 0.3340, Accuracy: 0.7891, TrueLabel: 0.5\n",
      "Step 2176, Loss: 0.3359, Accuracy: 0.7900, TrueLabel: 0.5\n",
      "Step 2208, Loss: 0.3281, Accuracy: 0.8013, TrueLabel: 0.5\n",
      "Step 2240, Loss: 0.3145, Accuracy: 0.7939, TrueLabel: 0.5\n",
      "Step 2272, Loss: 0.3281, Accuracy: 0.7876, TrueLabel: 0.5\n",
      "Step 2304, Loss: 0.3242, Accuracy: 0.7910, TrueLabel: 0.5\n",
      "Step 2336, Loss: 0.3262, Accuracy: 0.8052, TrueLabel: 0.5\n",
      "Step 2368, Loss: 0.3320, Accuracy: 0.7949, TrueLabel: 0.5\n",
      "Step 2400, Loss: 0.3223, Accuracy: 0.7939, TrueLabel: 0.5\n",
      "Step 2432, Loss: 0.3320, Accuracy: 0.7988, TrueLabel: 0.5\n",
      "Step 2464, Loss: 0.3203, Accuracy: 0.8008, TrueLabel: 0.5\n",
      "Step 2496, Loss: 0.3223, Accuracy: 0.8091, TrueLabel: 0.5\n",
      "Step 2528, Loss: 0.3223, Accuracy: 0.7915, TrueLabel: 0.5\n",
      "Step 2560, Loss: 0.3281, Accuracy: 0.7944, TrueLabel: 0.5\n",
      "Step 2592, Loss: 0.3223, Accuracy: 0.8013, TrueLabel: 0.5\n",
      "Step 2624, Loss: 0.3242, Accuracy: 0.8047, TrueLabel: 0.5\n",
      "Step 2656, Loss: 0.3398, Accuracy: 0.7891, TrueLabel: 0.5\n",
      "Step 2688, Loss: 0.3359, Accuracy: 0.8027, TrueLabel: 0.5\n",
      "Step 2720, Loss: 0.3281, Accuracy: 0.7832, TrueLabel: 0.5\n",
      "Step 2752, Loss: 0.3438, Accuracy: 0.7881, TrueLabel: 0.5\n",
      "Step 2784, Loss: 0.3203, Accuracy: 0.7837, TrueLabel: 0.5\n",
      "Step 2816, Loss: 0.3262, Accuracy: 0.7891, TrueLabel: 0.5\n",
      "Step 2848, Loss: 0.3418, Accuracy: 0.7832, TrueLabel: 0.5\n",
      "Step 2880, Loss: 0.3340, Accuracy: 0.7910, TrueLabel: 0.5\n",
      "Step 2912, Loss: 0.3262, Accuracy: 0.8008, TrueLabel: 0.5\n",
      "Step 2944, Loss: 0.3359, Accuracy: 0.7808, TrueLabel: 0.5\n",
      "Step 2976, Loss: 0.3262, Accuracy: 0.7954, TrueLabel: 0.5\n",
      "Step 3008, Loss: 0.3047, Accuracy: 0.8115, TrueLabel: 0.5\n",
      "Step 3040, Loss: 0.3203, Accuracy: 0.8037, TrueLabel: 0.5\n",
      "Step 3072, Loss: 0.3184, Accuracy: 0.8042, TrueLabel: 0.5\n",
      "Step 3104, Loss: 0.3203, Accuracy: 0.7944, TrueLabel: 0.5\n",
      "Step 3136, Loss: 0.3340, Accuracy: 0.7920, TrueLabel: 0.5\n",
      "Step 3168, Loss: 0.3066, Accuracy: 0.8032, TrueLabel: 0.5\n",
      "Step 3200, Loss: 0.3223, Accuracy: 0.7871, TrueLabel: 0.5\n",
      "Step 3232, Loss: 0.3223, Accuracy: 0.7959, TrueLabel: 0.5\n",
      "Step 3264, Loss: 0.3301, Accuracy: 0.7852, TrueLabel: 0.5\n",
      "Step 3296, Loss: 0.3320, Accuracy: 0.7925, TrueLabel: 0.5\n",
      "Step 3328, Loss: 0.3359, Accuracy: 0.7769, TrueLabel: 0.5\n",
      "Step 3360, Loss: 0.3105, Accuracy: 0.8120, TrueLabel: 0.5\n",
      "Step 3392, Loss: 0.3242, Accuracy: 0.7974, TrueLabel: 0.5\n",
      "Step 3424, Loss: 0.3398, Accuracy: 0.7822, TrueLabel: 0.5\n",
      "Step 3456, Loss: 0.3145, Accuracy: 0.8003, TrueLabel: 0.5\n",
      "Step 3488, Loss: 0.3125, Accuracy: 0.8120, TrueLabel: 0.5\n",
      "Step 3520, Loss: 0.3359, Accuracy: 0.7773, TrueLabel: 0.5\n",
      "Step 3552, Loss: 0.3262, Accuracy: 0.7900, TrueLabel: 0.5\n",
      "Step 3584, Loss: 0.3203, Accuracy: 0.8027, TrueLabel: 0.5\n",
      "Step 3616, Loss: 0.3203, Accuracy: 0.8003, TrueLabel: 0.5\n",
      "Step 3648, Loss: 0.3223, Accuracy: 0.7983, TrueLabel: 0.5\n",
      "Step 3680, Loss: 0.3125, Accuracy: 0.8022, TrueLabel: 0.5\n",
      "Step 3712, Loss: 0.3164, Accuracy: 0.7993, TrueLabel: 0.5\n",
      "Step 3744, Loss: 0.3320, Accuracy: 0.7871, TrueLabel: 0.5\n",
      "Step 3776, Loss: 0.3301, Accuracy: 0.7847, TrueLabel: 0.5\n",
      "Step 3808, Loss: 0.3203, Accuracy: 0.8071, TrueLabel: 0.5\n",
      "Step 3840, Loss: 0.3320, Accuracy: 0.7930, TrueLabel: 0.5\n",
      "Step 3872, Loss: 0.3203, Accuracy: 0.7847, TrueLabel: 0.5\n",
      "Step 3904, Loss: 0.3242, Accuracy: 0.7896, TrueLabel: 0.5\n",
      "Step 3936, Loss: 0.3301, Accuracy: 0.7827, TrueLabel: 0.5\n",
      "Step 3968, Loss: 0.3340, Accuracy: 0.7979, TrueLabel: 0.5\n",
      "Step 4000, Loss: 0.3262, Accuracy: 0.7910, TrueLabel: 0.5\n",
      "Step 4032, Loss: 0.3242, Accuracy: 0.7979, TrueLabel: 0.5\n",
      "Step 4064, Loss: 0.3242, Accuracy: 0.8076, TrueLabel: 0.5\n",
      "Step 4096, Loss: 0.3301, Accuracy: 0.7920, TrueLabel: 0.5\n",
      "Step 4128, Loss: 0.3320, Accuracy: 0.7993, TrueLabel: 0.5\n",
      "Step 4160, Loss: 0.3262, Accuracy: 0.7866, TrueLabel: 0.5\n",
      "Step 4192, Loss: 0.3281, Accuracy: 0.7759, TrueLabel: 0.5\n",
      "Step 4224, Loss: 0.3164, Accuracy: 0.7910, TrueLabel: 0.5\n",
      "Step 4256, Loss: 0.3203, Accuracy: 0.8096, TrueLabel: 0.5\n",
      "Step 4288, Loss: 0.3379, Accuracy: 0.7881, TrueLabel: 0.5\n",
      "Step 4320, Loss: 0.3359, Accuracy: 0.7900, TrueLabel: 0.5\n",
      "Step 4352, Loss: 0.3242, Accuracy: 0.7959, TrueLabel: 0.5\n",
      "Step 4384, Loss: 0.3105, Accuracy: 0.7993, TrueLabel: 0.5\n",
      "Step 4416, Loss: 0.3262, Accuracy: 0.7886, TrueLabel: 0.5\n",
      "Step 4448, Loss: 0.3203, Accuracy: 0.8076, TrueLabel: 0.5\n",
      "Step 4480, Loss: 0.3164, Accuracy: 0.8130, TrueLabel: 0.5\n",
      "Step 4512, Loss: 0.3438, Accuracy: 0.7749, TrueLabel: 0.5\n",
      "Step 4544, Loss: 0.3047, Accuracy: 0.8203, TrueLabel: 0.5\n",
      "Step 4576, Loss: 0.3320, Accuracy: 0.7852, TrueLabel: 0.5\n",
      "Step 4608, Loss: 0.3125, Accuracy: 0.8076, TrueLabel: 0.5\n",
      "Step 4640, Loss: 0.3262, Accuracy: 0.7817, TrueLabel: 0.5\n",
      "Step 4672, Loss: 0.3320, Accuracy: 0.7856, TrueLabel: 0.5\n",
      "Step 4704, Loss: 0.3184, Accuracy: 0.7988, TrueLabel: 0.5\n",
      "Step 4736, Loss: 0.3262, Accuracy: 0.8066, TrueLabel: 0.5\n",
      "Step 4768, Loss: 0.3184, Accuracy: 0.7974, TrueLabel: 0.5\n",
      "Step 4800, Loss: 0.3145, Accuracy: 0.8115, TrueLabel: 0.5\n",
      "Step 4832, Loss: 0.3164, Accuracy: 0.8057, TrueLabel: 0.5\n",
      "Step 4864, Loss: 0.3301, Accuracy: 0.7930, TrueLabel: 0.5\n",
      "Step 4896, Loss: 0.3320, Accuracy: 0.7930, TrueLabel: 0.5\n",
      "Step 4928, Loss: 0.3301, Accuracy: 0.7827, TrueLabel: 0.5\n",
      "Step 4960, Loss: 0.3301, Accuracy: 0.7988, TrueLabel: 0.5\n",
      "Step 4992, Loss: 0.3301, Accuracy: 0.7881, TrueLabel: 0.5\n",
      "\n",
      "=== Training with 4 layers ===\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "314f1f2286a245838a59fe963f8c13e1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0, Loss: 4.6562, Accuracy: 0.0005, TrueLabel: 0.5\n",
      "Step 32, Loss: 0.6484, Accuracy: 0.5806, TrueLabel: 0.5\n",
      "Step 64, Loss: 0.6211, Accuracy: 0.5952, TrueLabel: 0.5\n",
      "Step 96, Loss: 0.5664, Accuracy: 0.6392, TrueLabel: 0.5\n",
      "Step 128, Loss: 0.5469, Accuracy: 0.6436, TrueLabel: 0.5\n",
      "Step 160, Loss: 0.4609, Accuracy: 0.7153, TrueLabel: 0.5\n",
      "Step 192, Loss: 0.4473, Accuracy: 0.7100, TrueLabel: 0.5\n",
      "Step 224, Loss: 0.5195, Accuracy: 0.6475, TrueLabel: 0.5\n",
      "Step 256, Loss: 0.4043, Accuracy: 0.7500, TrueLabel: 0.5\n",
      "Step 288, Loss: 0.3887, Accuracy: 0.7451, TrueLabel: 0.5\n",
      "Step 320, Loss: 0.4453, Accuracy: 0.7163, TrueLabel: 0.5\n",
      "Step 352, Loss: 0.4316, Accuracy: 0.7271, TrueLabel: 0.5\n",
      "Step 384, Loss: 0.4316, Accuracy: 0.7188, TrueLabel: 0.5\n",
      "Step 416, Loss: 0.4082, Accuracy: 0.7295, TrueLabel: 0.5\n",
      "Step 448, Loss: 0.4824, Accuracy: 0.7085, TrueLabel: 0.5\n",
      "Step 480, Loss: 0.4395, Accuracy: 0.7026, TrueLabel: 0.5\n",
      "Step 512, Loss: 0.4043, Accuracy: 0.7461, TrueLabel: 0.5\n",
      "Step 544, Loss: 0.3984, Accuracy: 0.7446, TrueLabel: 0.5\n",
      "Step 576, Loss: 0.4004, Accuracy: 0.7363, TrueLabel: 0.5\n",
      "Step 608, Loss: 0.3945, Accuracy: 0.7529, TrueLabel: 0.5\n",
      "Step 640, Loss: 0.4141, Accuracy: 0.7407, TrueLabel: 0.5\n",
      "Step 672, Loss: 0.4062, Accuracy: 0.7480, TrueLabel: 0.5\n",
      "Step 704, Loss: 0.3730, Accuracy: 0.7583, TrueLabel: 0.5\n",
      "Step 736, Loss: 0.4258, Accuracy: 0.7339, TrueLabel: 0.5\n",
      "Step 768, Loss: 0.4043, Accuracy: 0.7466, TrueLabel: 0.5\n",
      "Step 800, Loss: 0.3965, Accuracy: 0.7422, TrueLabel: 0.5\n",
      "Step 832, Loss: 0.3926, Accuracy: 0.7539, TrueLabel: 0.5\n",
      "Step 864, Loss: 0.4336, Accuracy: 0.7495, TrueLabel: 0.5\n",
      "Step 896, Loss: 0.3965, Accuracy: 0.7485, TrueLabel: 0.5\n",
      "Step 928, Loss: 0.3945, Accuracy: 0.7632, TrueLabel: 0.5\n",
      "Step 960, Loss: 0.3770, Accuracy: 0.7769, TrueLabel: 0.5\n",
      "Step 992, Loss: 0.3848, Accuracy: 0.7505, TrueLabel: 0.5\n",
      "Step 1024, Loss: 0.3711, Accuracy: 0.7612, TrueLabel: 0.5\n",
      "Step 1056, Loss: 0.3594, Accuracy: 0.7817, TrueLabel: 0.5\n",
      "Step 1088, Loss: 0.3574, Accuracy: 0.7783, TrueLabel: 0.5\n",
      "Step 1120, Loss: 0.3379, Accuracy: 0.7861, TrueLabel: 0.5\n",
      "Step 1152, Loss: 0.3594, Accuracy: 0.7729, TrueLabel: 0.5\n",
      "Step 1184, Loss: 0.3301, Accuracy: 0.7900, TrueLabel: 0.5\n",
      "Step 1216, Loss: 0.3320, Accuracy: 0.7871, TrueLabel: 0.5\n",
      "Step 1248, Loss: 0.3379, Accuracy: 0.7861, TrueLabel: 0.5\n",
      "Step 1280, Loss: 0.3340, Accuracy: 0.7788, TrueLabel: 0.5\n",
      "Step 1312, Loss: 0.3281, Accuracy: 0.7783, TrueLabel: 0.5\n",
      "Step 1344, Loss: 0.3242, Accuracy: 0.7778, TrueLabel: 0.5\n",
      "Step 1376, Loss: 0.3477, Accuracy: 0.7949, TrueLabel: 0.5\n",
      "Step 1408, Loss: 0.3496, Accuracy: 0.7783, TrueLabel: 0.5\n",
      "Step 1440, Loss: 0.3281, Accuracy: 0.7856, TrueLabel: 0.5\n",
      "Step 1472, Loss: 0.3145, Accuracy: 0.7993, TrueLabel: 0.5\n",
      "Step 1504, Loss: 0.3848, Accuracy: 0.7598, TrueLabel: 0.5\n",
      "Step 1536, Loss: 0.3613, Accuracy: 0.7803, TrueLabel: 0.5\n",
      "Step 1568, Loss: 0.3340, Accuracy: 0.7939, TrueLabel: 0.5\n",
      "Step 1600, Loss: 0.3262, Accuracy: 0.7866, TrueLabel: 0.5\n",
      "Step 1632, Loss: 0.3262, Accuracy: 0.7886, TrueLabel: 0.5\n",
      "Step 1664, Loss: 0.3457, Accuracy: 0.7739, TrueLabel: 0.5\n",
      "Step 1696, Loss: 0.3184, Accuracy: 0.7925, TrueLabel: 0.5\n",
      "Step 1728, Loss: 0.3262, Accuracy: 0.8110, TrueLabel: 0.5\n",
      "Step 1760, Loss: 0.3262, Accuracy: 0.7876, TrueLabel: 0.5\n",
      "Step 1792, Loss: 0.3066, Accuracy: 0.8037, TrueLabel: 0.5\n",
      "Step 1824, Loss: 0.3477, Accuracy: 0.7759, TrueLabel: 0.5\n",
      "Step 1856, Loss: 0.3008, Accuracy: 0.8008, TrueLabel: 0.5\n",
      "Step 1888, Loss: 0.3301, Accuracy: 0.8008, TrueLabel: 0.5\n",
      "Step 1920, Loss: 0.3223, Accuracy: 0.8022, TrueLabel: 0.5\n",
      "Step 1952, Loss: 0.3262, Accuracy: 0.7915, TrueLabel: 0.5\n",
      "Step 1984, Loss: 0.3184, Accuracy: 0.7910, TrueLabel: 0.5\n",
      "Step 2016, Loss: 0.3145, Accuracy: 0.8008, TrueLabel: 0.5\n",
      "Step 2048, Loss: 0.3320, Accuracy: 0.7905, TrueLabel: 0.5\n",
      "Step 2080, Loss: 0.3242, Accuracy: 0.7998, TrueLabel: 0.5\n",
      "Step 2112, Loss: 0.3086, Accuracy: 0.8047, TrueLabel: 0.5\n",
      "Step 2144, Loss: 0.3125, Accuracy: 0.8057, TrueLabel: 0.5\n",
      "Step 2176, Loss: 0.3340, Accuracy: 0.7715, TrueLabel: 0.5\n",
      "Step 2208, Loss: 0.3027, Accuracy: 0.8145, TrueLabel: 0.5\n",
      "Step 2240, Loss: 0.3262, Accuracy: 0.7949, TrueLabel: 0.5\n",
      "Step 2272, Loss: 0.3125, Accuracy: 0.8066, TrueLabel: 0.5\n",
      "Step 2304, Loss: 0.3184, Accuracy: 0.7988, TrueLabel: 0.5\n",
      "Step 2336, Loss: 0.3281, Accuracy: 0.7969, TrueLabel: 0.5\n",
      "Step 2368, Loss: 0.3184, Accuracy: 0.7866, TrueLabel: 0.5\n",
      "Step 2400, Loss: 0.3164, Accuracy: 0.8013, TrueLabel: 0.5\n",
      "Step 2432, Loss: 0.3145, Accuracy: 0.8091, TrueLabel: 0.5\n",
      "Step 2464, Loss: 0.3203, Accuracy: 0.7959, TrueLabel: 0.5\n",
      "Step 2496, Loss: 0.3125, Accuracy: 0.7793, TrueLabel: 0.5\n",
      "Step 2528, Loss: 0.3145, Accuracy: 0.7944, TrueLabel: 0.5\n",
      "Step 2560, Loss: 0.2988, Accuracy: 0.8096, TrueLabel: 0.5\n",
      "Step 2592, Loss: 0.2988, Accuracy: 0.8062, TrueLabel: 0.5\n",
      "Step 2624, Loss: 0.3145, Accuracy: 0.7959, TrueLabel: 0.5\n",
      "Step 2656, Loss: 0.3008, Accuracy: 0.8149, TrueLabel: 0.5\n",
      "Step 2688, Loss: 0.3066, Accuracy: 0.8062, TrueLabel: 0.5\n",
      "Step 2720, Loss: 0.3008, Accuracy: 0.8042, TrueLabel: 0.5\n",
      "Step 2752, Loss: 0.2930, Accuracy: 0.8188, TrueLabel: 0.5\n",
      "Step 2784, Loss: 0.2949, Accuracy: 0.8140, TrueLabel: 0.5\n",
      "Step 2816, Loss: 0.2891, Accuracy: 0.8174, TrueLabel: 0.5\n",
      "Step 2848, Loss: 0.3145, Accuracy: 0.7871, TrueLabel: 0.5\n",
      "Step 2880, Loss: 0.3066, Accuracy: 0.8062, TrueLabel: 0.5\n",
      "Step 2912, Loss: 0.2969, Accuracy: 0.8086, TrueLabel: 0.5\n",
      "Step 2944, Loss: 0.3066, Accuracy: 0.8091, TrueLabel: 0.5\n",
      "Step 2976, Loss: 0.3047, Accuracy: 0.8042, TrueLabel: 0.5\n",
      "Step 3008, Loss: 0.2988, Accuracy: 0.8066, TrueLabel: 0.5\n",
      "Step 3040, Loss: 0.3184, Accuracy: 0.7949, TrueLabel: 0.5\n",
      "Step 3072, Loss: 0.3105, Accuracy: 0.8081, TrueLabel: 0.5\n",
      "Step 3104, Loss: 0.2949, Accuracy: 0.8110, TrueLabel: 0.5\n",
      "Step 3136, Loss: 0.3027, Accuracy: 0.8135, TrueLabel: 0.5\n",
      "Step 3168, Loss: 0.3164, Accuracy: 0.7905, TrueLabel: 0.5\n",
      "Step 3200, Loss: 0.3086, Accuracy: 0.8086, TrueLabel: 0.5\n",
      "Step 3232, Loss: 0.2930, Accuracy: 0.8091, TrueLabel: 0.5\n",
      "Step 3264, Loss: 0.3047, Accuracy: 0.8071, TrueLabel: 0.5\n",
      "Step 3296, Loss: 0.3066, Accuracy: 0.8120, TrueLabel: 0.5\n",
      "Step 3328, Loss: 0.3066, Accuracy: 0.8096, TrueLabel: 0.5\n",
      "Step 3360, Loss: 0.2930, Accuracy: 0.8096, TrueLabel: 0.5\n",
      "Step 3392, Loss: 0.2988, Accuracy: 0.8198, TrueLabel: 0.5\n",
      "Step 3424, Loss: 0.3047, Accuracy: 0.8096, TrueLabel: 0.5\n",
      "Step 3456, Loss: 0.2988, Accuracy: 0.8198, TrueLabel: 0.5\n",
      "Step 3488, Loss: 0.3008, Accuracy: 0.8032, TrueLabel: 0.5\n",
      "Step 3520, Loss: 0.2949, Accuracy: 0.8022, TrueLabel: 0.5\n",
      "Step 3552, Loss: 0.2910, Accuracy: 0.8193, TrueLabel: 0.5\n",
      "Step 3584, Loss: 0.2969, Accuracy: 0.8052, TrueLabel: 0.5\n",
      "Step 3616, Loss: 0.3008, Accuracy: 0.8076, TrueLabel: 0.5\n",
      "Step 3648, Loss: 0.2871, Accuracy: 0.8223, TrueLabel: 0.5\n",
      "Step 3680, Loss: 0.3008, Accuracy: 0.8164, TrueLabel: 0.5\n",
      "Step 3712, Loss: 0.3164, Accuracy: 0.8062, TrueLabel: 0.5\n",
      "Step 3744, Loss: 0.2930, Accuracy: 0.8179, TrueLabel: 0.5\n",
      "Step 3776, Loss: 0.3086, Accuracy: 0.8042, TrueLabel: 0.5\n",
      "Step 3808, Loss: 0.2930, Accuracy: 0.8174, TrueLabel: 0.5\n",
      "Step 3840, Loss: 0.2988, Accuracy: 0.8242, TrueLabel: 0.5\n",
      "Step 3872, Loss: 0.2949, Accuracy: 0.8154, TrueLabel: 0.5\n",
      "Step 3904, Loss: 0.2988, Accuracy: 0.8110, TrueLabel: 0.5\n",
      "Step 3936, Loss: 0.2852, Accuracy: 0.8276, TrueLabel: 0.5\n",
      "Step 3968, Loss: 0.2930, Accuracy: 0.8062, TrueLabel: 0.5\n",
      "Step 4000, Loss: 0.2988, Accuracy: 0.8208, TrueLabel: 0.5\n",
      "Step 4032, Loss: 0.2832, Accuracy: 0.8262, TrueLabel: 0.5\n",
      "Step 4064, Loss: 0.2910, Accuracy: 0.8198, TrueLabel: 0.5\n",
      "Step 4096, Loss: 0.2969, Accuracy: 0.8169, TrueLabel: 0.5\n",
      "Step 4128, Loss: 0.2988, Accuracy: 0.8174, TrueLabel: 0.5\n",
      "Step 4160, Loss: 0.3047, Accuracy: 0.8081, TrueLabel: 0.5\n",
      "Step 4192, Loss: 0.2852, Accuracy: 0.8379, TrueLabel: 0.5\n",
      "Step 4224, Loss: 0.2891, Accuracy: 0.8213, TrueLabel: 0.5\n",
      "Step 4256, Loss: 0.2949, Accuracy: 0.8120, TrueLabel: 0.5\n",
      "Step 4288, Loss: 0.2715, Accuracy: 0.8301, TrueLabel: 0.5\n",
      "Step 4320, Loss: 0.3066, Accuracy: 0.8047, TrueLabel: 0.5\n",
      "Step 4352, Loss: 0.2949, Accuracy: 0.8198, TrueLabel: 0.5\n",
      "Step 4384, Loss: 0.2949, Accuracy: 0.8120, TrueLabel: 0.5\n",
      "Step 4416, Loss: 0.2891, Accuracy: 0.8164, TrueLabel: 0.5\n",
      "Step 4448, Loss: 0.2969, Accuracy: 0.8198, TrueLabel: 0.5\n",
      "Step 4480, Loss: 0.2871, Accuracy: 0.8179, TrueLabel: 0.5\n",
      "Step 4512, Loss: 0.3086, Accuracy: 0.8018, TrueLabel: 0.5\n",
      "Step 4544, Loss: 0.2832, Accuracy: 0.8242, TrueLabel: 0.5\n",
      "Step 4576, Loss: 0.2852, Accuracy: 0.8193, TrueLabel: 0.5\n",
      "Step 4608, Loss: 0.2754, Accuracy: 0.8306, TrueLabel: 0.5\n",
      "Step 4640, Loss: 0.2812, Accuracy: 0.8140, TrueLabel: 0.5\n",
      "Step 4672, Loss: 0.3047, Accuracy: 0.8159, TrueLabel: 0.5\n",
      "Step 4704, Loss: 0.3008, Accuracy: 0.8184, TrueLabel: 0.5\n",
      "Step 4736, Loss: 0.2969, Accuracy: 0.8179, TrueLabel: 0.5\n",
      "Step 4768, Loss: 0.2969, Accuracy: 0.8081, TrueLabel: 0.5\n",
      "Step 4800, Loss: 0.3027, Accuracy: 0.7979, TrueLabel: 0.5\n",
      "Step 4832, Loss: 0.2832, Accuracy: 0.8276, TrueLabel: 0.5\n",
      "Step 4864, Loss: 0.2988, Accuracy: 0.8140, TrueLabel: 0.5\n",
      "Step 4896, Loss: 0.2930, Accuracy: 0.8013, TrueLabel: 0.5\n",
      "Step 4928, Loss: 0.2988, Accuracy: 0.8140, TrueLabel: 0.5\n",
      "Step 4960, Loss: 0.2988, Accuracy: 0.8262, TrueLabel: 0.5\n",
      "Step 4992, Loss: 0.2910, Accuracy: 0.8110, TrueLabel: 0.5\n",
      "\n",
      "=== Training with 8 layers ===\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9fda782b17d349bf81544b74eddf2831",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0, Loss: 5.4375, Accuracy: 0.0000, TrueLabel: 0.5\n",
      "Step 32, Loss: 0.6523, Accuracy: 0.5767, TrueLabel: 0.5\n",
      "Step 64, Loss: 0.7227, Accuracy: 0.5522, TrueLabel: 0.5\n",
      "Step 96, Loss: 0.5977, Accuracy: 0.6133, TrueLabel: 0.5\n",
      "Step 128, Loss: 0.5312, Accuracy: 0.6606, TrueLabel: 0.5\n",
      "Step 160, Loss: 0.5078, Accuracy: 0.6758, TrueLabel: 0.5\n",
      "Step 192, Loss: 0.4805, Accuracy: 0.6821, TrueLabel: 0.5\n",
      "Step 224, Loss: 0.4707, Accuracy: 0.6895, TrueLabel: 0.5\n",
      "Step 256, Loss: 0.4355, Accuracy: 0.7188, TrueLabel: 0.5\n",
      "Step 288, Loss: 0.4180, Accuracy: 0.7344, TrueLabel: 0.5\n",
      "Step 320, Loss: 0.4492, Accuracy: 0.6924, TrueLabel: 0.5\n",
      "Step 352, Loss: 0.4648, Accuracy: 0.6924, TrueLabel: 0.5\n",
      "Step 384, Loss: 0.4199, Accuracy: 0.7261, TrueLabel: 0.5\n",
      "Step 416, Loss: 0.4121, Accuracy: 0.7358, TrueLabel: 0.5\n",
      "Step 448, Loss: 0.4062, Accuracy: 0.7402, TrueLabel: 0.5\n",
      "Step 480, Loss: 0.3906, Accuracy: 0.7529, TrueLabel: 0.5\n",
      "Step 512, Loss: 0.3945, Accuracy: 0.7441, TrueLabel: 0.5\n",
      "Step 544, Loss: 0.4121, Accuracy: 0.7383, TrueLabel: 0.5\n",
      "Step 576, Loss: 0.3926, Accuracy: 0.7578, TrueLabel: 0.5\n",
      "Step 608, Loss: 0.4160, Accuracy: 0.7417, TrueLabel: 0.5\n",
      "Step 640, Loss: 0.3887, Accuracy: 0.7461, TrueLabel: 0.5\n",
      "Step 672, Loss: 0.4004, Accuracy: 0.7271, TrueLabel: 0.5\n",
      "Step 704, Loss: 0.3926, Accuracy: 0.7495, TrueLabel: 0.5\n",
      "Step 736, Loss: 0.4082, Accuracy: 0.7388, TrueLabel: 0.5\n",
      "Step 768, Loss: 0.3809, Accuracy: 0.7749, TrueLabel: 0.5\n",
      "Step 800, Loss: 0.4141, Accuracy: 0.7246, TrueLabel: 0.5\n",
      "Step 832, Loss: 0.3730, Accuracy: 0.7754, TrueLabel: 0.5\n",
      "Step 864, Loss: 0.3574, Accuracy: 0.7651, TrueLabel: 0.5\n",
      "Step 896, Loss: 0.3516, Accuracy: 0.7661, TrueLabel: 0.5\n",
      "Step 928, Loss: 0.3965, Accuracy: 0.7539, TrueLabel: 0.5\n",
      "Step 960, Loss: 0.3574, Accuracy: 0.7754, TrueLabel: 0.5\n",
      "Step 992, Loss: 0.3594, Accuracy: 0.7676, TrueLabel: 0.5\n",
      "Step 1024, Loss: 0.3711, Accuracy: 0.7480, TrueLabel: 0.5\n",
      "Step 1056, Loss: 0.3691, Accuracy: 0.7725, TrueLabel: 0.5\n",
      "Step 1088, Loss: 0.3555, Accuracy: 0.7744, TrueLabel: 0.5\n",
      "Step 1120, Loss: 0.3555, Accuracy: 0.7705, TrueLabel: 0.5\n",
      "Step 1152, Loss: 0.3379, Accuracy: 0.7832, TrueLabel: 0.5\n",
      "Step 1184, Loss: 0.3828, Accuracy: 0.7559, TrueLabel: 0.5\n",
      "Step 1216, Loss: 0.3613, Accuracy: 0.7646, TrueLabel: 0.5\n",
      "Step 1248, Loss: 0.3457, Accuracy: 0.7866, TrueLabel: 0.5\n",
      "Step 1280, Loss: 0.3555, Accuracy: 0.7720, TrueLabel: 0.5\n",
      "Step 1312, Loss: 0.3711, Accuracy: 0.7671, TrueLabel: 0.5\n",
      "Step 1344, Loss: 0.3633, Accuracy: 0.7798, TrueLabel: 0.5\n",
      "Step 1376, Loss: 0.3711, Accuracy: 0.7808, TrueLabel: 0.5\n",
      "Step 1408, Loss: 0.3652, Accuracy: 0.7700, TrueLabel: 0.5\n",
      "Step 1440, Loss: 0.3594, Accuracy: 0.7749, TrueLabel: 0.5\n",
      "Step 1472, Loss: 0.3516, Accuracy: 0.7725, TrueLabel: 0.5\n",
      "Step 1504, Loss: 0.3477, Accuracy: 0.7690, TrueLabel: 0.5\n",
      "Step 1536, Loss: 0.3633, Accuracy: 0.7622, TrueLabel: 0.5\n",
      "Step 1568, Loss: 0.3496, Accuracy: 0.7822, TrueLabel: 0.5\n",
      "Step 1600, Loss: 0.3652, Accuracy: 0.7646, TrueLabel: 0.5\n",
      "Step 1632, Loss: 0.3672, Accuracy: 0.7539, TrueLabel: 0.5\n",
      "Step 1664, Loss: 0.3613, Accuracy: 0.7700, TrueLabel: 0.5\n",
      "Step 1696, Loss: 0.3711, Accuracy: 0.7563, TrueLabel: 0.5\n",
      "Step 1728, Loss: 0.3340, Accuracy: 0.7852, TrueLabel: 0.5\n",
      "Step 1760, Loss: 0.3535, Accuracy: 0.7793, TrueLabel: 0.5\n",
      "Step 1792, Loss: 0.3613, Accuracy: 0.7666, TrueLabel: 0.5\n",
      "Step 1824, Loss: 0.3398, Accuracy: 0.7803, TrueLabel: 0.5\n",
      "Step 1856, Loss: 0.3398, Accuracy: 0.7944, TrueLabel: 0.5\n",
      "Step 1888, Loss: 0.3418, Accuracy: 0.7734, TrueLabel: 0.5\n",
      "Step 1920, Loss: 0.3965, Accuracy: 0.7500, TrueLabel: 0.5\n",
      "Step 1952, Loss: 0.4062, Accuracy: 0.7607, TrueLabel: 0.5\n",
      "Step 1984, Loss: 0.3594, Accuracy: 0.7695, TrueLabel: 0.5\n",
      "Step 2016, Loss: 0.3770, Accuracy: 0.7627, TrueLabel: 0.5\n",
      "Step 2048, Loss: 0.3398, Accuracy: 0.7783, TrueLabel: 0.5\n",
      "Step 2080, Loss: 0.3574, Accuracy: 0.7466, TrueLabel: 0.5\n",
      "Step 2112, Loss: 0.3398, Accuracy: 0.7939, TrueLabel: 0.5\n",
      "Step 2144, Loss: 0.3262, Accuracy: 0.7964, TrueLabel: 0.5\n",
      "Step 2176, Loss: 0.3242, Accuracy: 0.8057, TrueLabel: 0.5\n",
      "Step 2208, Loss: 0.3301, Accuracy: 0.7974, TrueLabel: 0.5\n",
      "Step 2240, Loss: 0.3301, Accuracy: 0.7896, TrueLabel: 0.5\n",
      "Step 2272, Loss: 0.3379, Accuracy: 0.7886, TrueLabel: 0.5\n",
      "Step 2304, Loss: 0.3223, Accuracy: 0.7964, TrueLabel: 0.5\n",
      "Step 2336, Loss: 0.3203, Accuracy: 0.7930, TrueLabel: 0.5\n",
      "Step 2368, Loss: 0.3340, Accuracy: 0.7817, TrueLabel: 0.5\n",
      "Step 2400, Loss: 0.3223, Accuracy: 0.7847, TrueLabel: 0.5\n",
      "Step 2432, Loss: 0.3184, Accuracy: 0.8145, TrueLabel: 0.5\n",
      "Step 2464, Loss: 0.3398, Accuracy: 0.7939, TrueLabel: 0.5\n",
      "Step 2496, Loss: 0.3164, Accuracy: 0.7871, TrueLabel: 0.5\n",
      "Step 2528, Loss: 0.3301, Accuracy: 0.7949, TrueLabel: 0.5\n",
      "Step 2560, Loss: 0.3242, Accuracy: 0.7910, TrueLabel: 0.5\n",
      "Step 2592, Loss: 0.3555, Accuracy: 0.7642, TrueLabel: 0.5\n",
      "Step 2624, Loss: 0.3340, Accuracy: 0.7930, TrueLabel: 0.5\n",
      "Step 2656, Loss: 0.3398, Accuracy: 0.7808, TrueLabel: 0.5\n",
      "Step 2688, Loss: 0.3477, Accuracy: 0.7656, TrueLabel: 0.5\n",
      "Step 2720, Loss: 0.3242, Accuracy: 0.7959, TrueLabel: 0.5\n",
      "Step 2752, Loss: 0.3574, Accuracy: 0.7734, TrueLabel: 0.5\n",
      "Step 2784, Loss: 0.3320, Accuracy: 0.7896, TrueLabel: 0.5\n",
      "Step 2816, Loss: 0.3223, Accuracy: 0.7979, TrueLabel: 0.5\n",
      "Step 2848, Loss: 0.3223, Accuracy: 0.7905, TrueLabel: 0.5\n",
      "Step 2880, Loss: 0.3281, Accuracy: 0.7935, TrueLabel: 0.5\n",
      "Step 2912, Loss: 0.3242, Accuracy: 0.8018, TrueLabel: 0.5\n",
      "Step 2944, Loss: 0.3281, Accuracy: 0.8105, TrueLabel: 0.5\n",
      "Step 2976, Loss: 0.3203, Accuracy: 0.7852, TrueLabel: 0.5\n",
      "Step 3008, Loss: 0.3301, Accuracy: 0.7891, TrueLabel: 0.5\n",
      "Step 3040, Loss: 0.3359, Accuracy: 0.7793, TrueLabel: 0.5\n",
      "Step 3072, Loss: 0.3301, Accuracy: 0.7705, TrueLabel: 0.5\n",
      "Step 3104, Loss: 0.3438, Accuracy: 0.7817, TrueLabel: 0.5\n",
      "Step 3136, Loss: 0.3105, Accuracy: 0.8037, TrueLabel: 0.5\n",
      "Step 3168, Loss: 0.3184, Accuracy: 0.8110, TrueLabel: 0.5\n",
      "Step 3200, Loss: 0.3379, Accuracy: 0.7793, TrueLabel: 0.5\n",
      "Step 3232, Loss: 0.3145, Accuracy: 0.8096, TrueLabel: 0.5\n",
      "Step 3264, Loss: 0.3262, Accuracy: 0.8003, TrueLabel: 0.5\n",
      "Step 3296, Loss: 0.3086, Accuracy: 0.8110, TrueLabel: 0.5\n",
      "Step 3328, Loss: 0.3457, Accuracy: 0.7759, TrueLabel: 0.5\n",
      "Step 3360, Loss: 0.3203, Accuracy: 0.7998, TrueLabel: 0.5\n",
      "Step 3392, Loss: 0.3145, Accuracy: 0.7988, TrueLabel: 0.5\n",
      "Step 3424, Loss: 0.3145, Accuracy: 0.7993, TrueLabel: 0.5\n",
      "Step 3456, Loss: 0.3125, Accuracy: 0.7964, TrueLabel: 0.5\n",
      "Step 3488, Loss: 0.3027, Accuracy: 0.8145, TrueLabel: 0.5\n",
      "Step 3520, Loss: 0.3105, Accuracy: 0.8018, TrueLabel: 0.5\n",
      "Step 3552, Loss: 0.3164, Accuracy: 0.7944, TrueLabel: 0.5\n",
      "Step 3584, Loss: 0.3125, Accuracy: 0.8130, TrueLabel: 0.5\n",
      "Step 3616, Loss: 0.3203, Accuracy: 0.7886, TrueLabel: 0.5\n",
      "Step 3648, Loss: 0.3145, Accuracy: 0.7935, TrueLabel: 0.5\n",
      "Step 3680, Loss: 0.3027, Accuracy: 0.7959, TrueLabel: 0.5\n",
      "Step 3712, Loss: 0.3047, Accuracy: 0.7993, TrueLabel: 0.5\n",
      "Step 3744, Loss: 0.3066, Accuracy: 0.8091, TrueLabel: 0.5\n",
      "Step 3776, Loss: 0.3086, Accuracy: 0.8057, TrueLabel: 0.5\n",
      "Step 3808, Loss: 0.3066, Accuracy: 0.8076, TrueLabel: 0.5\n",
      "Step 3840, Loss: 0.3184, Accuracy: 0.7876, TrueLabel: 0.5\n",
      "Step 3872, Loss: 0.3203, Accuracy: 0.7871, TrueLabel: 0.5\n",
      "Step 3904, Loss: 0.2930, Accuracy: 0.8267, TrueLabel: 0.5\n",
      "Step 3936, Loss: 0.2988, Accuracy: 0.8057, TrueLabel: 0.5\n",
      "Step 3968, Loss: 0.2891, Accuracy: 0.8271, TrueLabel: 0.5\n",
      "Step 4000, Loss: 0.3125, Accuracy: 0.8022, TrueLabel: 0.5\n",
      "Step 4032, Loss: 0.3047, Accuracy: 0.7969, TrueLabel: 0.5\n",
      "Step 4064, Loss: 0.2949, Accuracy: 0.8208, TrueLabel: 0.5\n",
      "Step 4096, Loss: 0.3066, Accuracy: 0.8105, TrueLabel: 0.5\n",
      "Step 4128, Loss: 0.2988, Accuracy: 0.8076, TrueLabel: 0.5\n",
      "Step 4160, Loss: 0.3047, Accuracy: 0.8027, TrueLabel: 0.5\n",
      "Step 4192, Loss: 0.2832, Accuracy: 0.8257, TrueLabel: 0.5\n",
      "Step 4224, Loss: 0.2969, Accuracy: 0.8062, TrueLabel: 0.5\n",
      "Step 4256, Loss: 0.3125, Accuracy: 0.8062, TrueLabel: 0.5\n",
      "Step 4288, Loss: 0.2988, Accuracy: 0.8101, TrueLabel: 0.5\n",
      "Step 4320, Loss: 0.2891, Accuracy: 0.8115, TrueLabel: 0.5\n",
      "Step 4352, Loss: 0.3047, Accuracy: 0.7993, TrueLabel: 0.5\n",
      "Step 4384, Loss: 0.2891, Accuracy: 0.8281, TrueLabel: 0.5\n",
      "Step 4416, Loss: 0.2988, Accuracy: 0.8135, TrueLabel: 0.5\n",
      "Step 4448, Loss: 0.3027, Accuracy: 0.8125, TrueLabel: 0.5\n",
      "Step 4480, Loss: 0.3262, Accuracy: 0.7964, TrueLabel: 0.5\n",
      "Step 4512, Loss: 0.3047, Accuracy: 0.7949, TrueLabel: 0.5\n",
      "Step 4544, Loss: 0.3027, Accuracy: 0.8120, TrueLabel: 0.5\n",
      "Step 4576, Loss: 0.3047, Accuracy: 0.7949, TrueLabel: 0.5\n",
      "Step 4608, Loss: 0.3164, Accuracy: 0.7988, TrueLabel: 0.5\n",
      "Step 4640, Loss: 0.2988, Accuracy: 0.8193, TrueLabel: 0.5\n",
      "Step 4672, Loss: 0.2949, Accuracy: 0.8101, TrueLabel: 0.5\n",
      "Step 4704, Loss: 0.3203, Accuracy: 0.7900, TrueLabel: 0.5\n",
      "Step 4736, Loss: 0.2988, Accuracy: 0.7993, TrueLabel: 0.5\n",
      "Step 4768, Loss: 0.2949, Accuracy: 0.8164, TrueLabel: 0.5\n",
      "Step 4800, Loss: 0.3008, Accuracy: 0.8052, TrueLabel: 0.5\n",
      "Step 4832, Loss: 0.3027, Accuracy: 0.8018, TrueLabel: 0.5\n",
      "Step 4864, Loss: 0.3086, Accuracy: 0.7954, TrueLabel: 0.5\n",
      "Step 4896, Loss: 0.3066, Accuracy: 0.8057, TrueLabel: 0.5\n",
      "Step 4928, Loss: 0.2852, Accuracy: 0.8286, TrueLabel: 0.5\n",
      "Step 4960, Loss: 0.3008, Accuracy: 0.8086, TrueLabel: 0.5\n",
      "Step 4992, Loss: 0.2871, Accuracy: 0.8223, TrueLabel: 0.5\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 85348 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import AdamW, get_scheduler\n",
    "from tqdm.auto import tqdm\n",
    "import fla\n",
    "from fla.models.utils import Cache\n",
    "from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss\n",
    "from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast\n",
    "from transformers import AutoModelForCausalLM\n",
    "import random\n",
    "from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union\n",
    "from itertools import permutations\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "def create_graph(n):\n",
    "    if n % 2 != 0:\n",
    "        raise ValueError(\"n should be an even number\")\n",
    "    \n",
    "    # Step 1: Randomly divide the points into two sets S_1 and S_2\n",
    "    points = list(range(1, n + 1))\n",
    "    random.shuffle(points)\n",
    "    mid = n // 2\n",
    "    S_1, S_2 = sorted(points[:mid]), sorted(points[mid:])\n",
    "    \n",
    "    def assign_parents(S):\n",
    "        parents = {}\n",
    "        for i in range(1, len(S)):\n",
    "            possible_parents = S[:i]\n",
    "            parents[S[i]] = random.choice(possible_parents) \n",
    "        return parents\n",
    "    \n",
    "    # Step 2: Assign parent nodes within each set\n",
    "    parents_S1 = assign_parents(S_1)\n",
    "    parents_S2 = assign_parents(S_2)\n",
    "    \n",
    "    # Step 3: Build adjacency matrix\n",
    "    adjacency_matrix = np.eye(n)\n",
    "    def fill_adjacency_matrix(parents):\n",
    "        for child, parent in parents.items():\n",
    "            if parent is not None:\n",
    "                adjacency_matrix[child - 1][parent - 1] = 1\n",
    "\n",
    "    fill_adjacency_matrix(parents_S1)\n",
    "    fill_adjacency_matrix(parents_S2)\n",
    "    labels = [0 for i in range(n)]\n",
    "    if 1 in S_1:\n",
    "        for i in S_1:\n",
    "            labels[i-1] = 1\n",
    "    else:\n",
    "        for i in S_2:\n",
    "            labels[i-1] = 1\n",
    "    return labels, adjacency_matrix\n",
    "\n",
    "def generate_graph_data(num_samples=100, n=32):\n",
    "    \"\"\"\n",
    "    Generates graph data samples with reachability information.\n",
    "    \n",
    "    :param num_samples: Number of samples to generate.\n",
    "    :param n: Number of nodes in the graph.\n",
    "    :return: A list of tuples. Each tuple contains an adjacency matrix and a list of labels indicating reachability from node 1 to each node.\n",
    "    \"\"\"\n",
    "    data = []\n",
    "\n",
    "    for _ in range(num_samples):\n",
    "        labels, A = create_graph(n)\n",
    "\n",
    "        # 将邻接矩阵和label转换为张量\n",
    "        adj_matrix = torch.tensor(A, dtype=torch.float)\n",
    "        # adj_matrix  =  adj_matrix.transpose(0,1)\n",
    "\n",
    "        labels = torch.tensor(labels, dtype=torch.long)\n",
    "\n",
    "        # 添加到数据中\n",
    "        data.append((adj_matrix, labels))\n",
    "    return data\n",
    "\n",
    "# 自定义 Dataset 类\n",
    "class S5Dataset(Dataset):\n",
    "    def __init__(self, k=4, num_samples=3):\n",
    "        self.k = k\n",
    "        self.num_samples = num_samples\n",
    "        self.data = generate_graph_data( self.num_samples,self.k)\n",
    "    def __len__(self):\n",
    "        return self.num_samples  # 每次只返回 num_samples 条数据\n",
    "    def data_reset(self):\n",
    "        self.data = generate_graph_data( self.num_samples,self.k)\n",
    "    def __getitem__(self, idx):\n",
    "        # 每次访问时重新生成数据，保证每个 epoch 数据不同\n",
    "        \n",
    "        input_ids, labels = self.data[idx]\n",
    "        return {\n",
    "            \"input_ids\": input_ids,\n",
    "            \"labels\": labels\n",
    "        }\n",
    "\n",
    "\n",
    "dtype = torch.bfloat16\n",
    "device = \"cuda:1\" if torch.cuda.is_available() else \"cpu\"\n",
    "class Emb(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.hidden_size = config.hidden_size\n",
    "    def forward(self, x):\n",
    "        # x shape: (batch_size, seq_len, input_dim)\n",
    "        batch_size, seq_len, input_dim = x.shape\n",
    "\n",
    "        # Step 1: 构建 one-hot 形式的位置编码（只在第一个位置加）\n",
    "        # 创建一个长度为 seq_len 的 one-hot 向量，只有第一个位置是 1\n",
    "        pos_onehot = torch.zeros(seq_len, seq_len, device=x.device)\n",
    "        pos_onehot[0, 0] = 1  # 第一个位置为 one-hot\n",
    "        pos_emb = pos_onehot.unsqueeze(0).expand(batch_size, -1, -1)  # (batch_size, seq_len, seq_len)\n",
    "\n",
    "        # # Step 2: 拼接到输入 x 的最后一个维度\n",
    "        # x = torch.cat([x, pos_emb], dim=-1)  # (batch_size, seq_len, input_dim + seq_len)\n",
    "\n",
    "        # Step 3: 如果拼接后仍小于 hidden_size，则继续用 zero padding 补齐\n",
    "        current_dim = x.size(-1)\n",
    "        if current_dim < self.hidden_size:\n",
    "            pad_size = list(x.shape)\n",
    "            pad_size[-1] = self.hidden_size - current_dim\n",
    "            padding = torch.zeros(*pad_size, device=x.device)\n",
    "            x = torch.cat([x, padding], dim=-1)  # (batch_size, seq_len, hidden_size)\n",
    "        else:\n",
    "            assert False\n",
    "        return x.to(dtype) \n",
    "        return torch.nn.functional.normalize(x+1e-5,dim=-1).to(dtype)\n",
    "    \n",
    "class Iden(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# 训练函数\n",
    "def train_model_with_layers(num_layers, run_name=\"run\", log_dict=None,length = 32):\n",
    "# 参数设置\n",
    "    # config = fla.models.DeltaTransformerConfig(\n",
    "    #         vocab_size=128,\n",
    "    #         hidden_size=1024,\n",
    "    #         num_hidden_layers=num_layers,\n",
    "    #         num_heads=8,\n",
    "    #         num_kv_heads=8,\n",
    "    #         max_position_embeddings=2048,\n",
    "    #         layer_norm_eps=1e-6,\n",
    "    #         # use_rotary_pos_emb=None,\n",
    "    #         fuse_cross_entropy = False,\n",
    "    #         rope_theta = -1,\n",
    "    #         kl2 = False,\n",
    "    #         kernel = 'round',\n",
    "    #         usequ = True,\n",
    "    #         bos_token_id = 127,\n",
    "    #         eos_token_id = 126,\n",
    "    #         returnu = True, \n",
    "    #         fuse_norm= False,\n",
    "    #         fuse_swiglu = False,\n",
    "    #     )\n",
    "\n",
    "    config = fla.models.TransformerConfig(\n",
    "        vocab_size=128,\n",
    "        hidden_size=1024,\n",
    "        num_hidden_layers=num_layers,\n",
    "        num_heads=8,\n",
    "        num_kv_heads=8,\n",
    "        max_position_embeddings=2048,\n",
    "        layer_norm_eps=1e-6,\n",
    "        use_rotary_pos_emb=None,\n",
    "        fuse_cross_entropy = False,\n",
    "        rope_theta = 1e4,\n",
    "    )\n",
    "    model = AutoModelForCausalLM.from_config(config)\n",
    "    model.model.embeddings = Emb(config)    \n",
    "    # model.model.layers[0].mlp = Iden(config)\n",
    "    model.to(device).to(dtype)\n",
    "\n",
    "    # 超参数\n",
    "    num_epochs = 5000\n",
    "    batch_size = 32\n",
    "    k = int(length)\n",
    "    num_samples_per_batch = batch_size  # 每次生成的样本数 = batch size\n",
    "\n",
    "    # 数据集和数据加载器\n",
    "    dataset = S5Dataset(k=k, num_samples=num_samples_per_batch)\n",
    "    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    # 优化器和学习率调度器\n",
    "    optimizer = AdamW(model.parameters(), lr=1e-3,weight_decay=0.1)\n",
    "    num_training_steps = num_epochs * len(dataloader)\n",
    "    lr_scheduler = get_scheduler(\n",
    "        name=\"linear\", optimizer=optimizer, num_warmup_steps=128, num_training_steps=num_training_steps\n",
    "    )\n",
    "\n",
    "    # 训练循环\n",
    "    progress_bar = tqdm(range(num_training_steps))\n",
    "\n",
    "    model.train()\n",
    "    step = 0\n",
    "    acc_log = []\n",
    "    loss_log = []\n",
    "    for epoch in range(num_epochs):\n",
    "        dataloader.dataset.data_reset()\n",
    "        for batch in dataloader:\n",
    "            input_ids = batch[\"input_ids\"].to(device)\n",
    "            labels = batch[\"labels\"].to(device)\n",
    "            outputs = model(\n",
    "                input_ids=input_ids,\n",
    "                labels=labels,\n",
    "                return_dict = True\n",
    "            )\n",
    "\n",
    "            loss = outputs.loss\n",
    "            # 准确率计算\n",
    "            logits = outputs.logits\n",
    "            # print(outputs)\n",
    "            # print(input_ids[0])\n",
    "            # print(labels[0])\n",
    "            # assert 1==2\n",
    "            predictions = torch.argmax(logits, dim=-1)\n",
    "            correct = (predictions == labels).sum().item()\n",
    "\n",
    "            \n",
    "            total = labels.numel()\n",
    "            accuracy = correct / total\n",
    "\n",
    "            # 反向传播\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # 打印信息\n",
    "            if step%32==0:\n",
    "                num = labels.float().mean().item()\n",
    "                print(f\"Step {step}, Loss: {loss.item():.4f}, Accuracy: {accuracy:.4f}, TrueLabel: {num}\")\n",
    "            step += 1\n",
    "\n",
    "            progress_bar.update(1)\n",
    "            acc_log.append(accuracy)\n",
    "            loss_log.append(loss.item())\n",
    "        # if accuracy>0.99:\n",
    "        #     k*=2\n",
    "        #     print(f'setp {step}, length come to {k}')\n",
    "        #     dataset = S5Dataset(k=k, num_samples=num_samples_per_batch)\n",
    "        #     dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "    log_dict[run_name] = {\"acc\": acc_log, \"loss\": loss_log}\n",
    "if __name__ == \"__main__\":\n",
    "    log_dict = {}\n",
    "\n",
    "    # 不同层数跑实验\n",
    "    for layers in [1,2,4,8]:\n",
    "        print(f\"\\n=== Training with {layers} layers ===\")\n",
    "        train_model_with_layers(layers, run_name=f\"{layers}_layers\", log_dict=log_dict,length=64)\n",
    "\n",
    "    # 绘图\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    for key in log_dict:\n",
    "        accs = log_dict[key]['acc']\n",
    "        steps = np.arange(len(accs))\n",
    "        plt.plot(steps, accs, label=key)\n",
    "\n",
    "    # plt.title(\"Accuracy vs Training Steps for Different Layer Counts\")\n",
    "    plt.xlabel(\"Training Step\")\n",
    "    plt.ylabel(\"Accuracy\")\n",
    "\n",
    "    plt.ylim(0.5,1)\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.savefig(\"accuracy_comparison.pdf\")\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "fileId": "bfa8093b-5162-4258-99d5-20d380b72f5f",
  "filePath": "/mnt/bn/seed-aws-va/xumingyu/workspaces/ARC-AGI/tetttt.ipynb",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
