learn_all_bits.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef8eb8e5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import sympy\n",
    "import numpy as np\n",
    "import random\n",
    "from tqdm import tqdm\n",
    "from fractions import Fraction\n",
    "\n",
    "bits = 40\n",
    "#generate prime number with given bit length\n",
    "flag = True\n",
    "while flag:\n",
    "    b=random.getrandbits(bits)\n",
    "    p = sympy.nextprime(b)\n",
    "    bin_p = bin(p)[2:]\n",
    "    if len(bin_p)==bits:\n",
    "        flag=False\n",
    "max_len = len(bin_p)\n",
    "#select a\n",
    "below=0\n",
    "while(below==0):\n",
    "    a = random.getrandbits(max_len)\n",
    "    if(a<=p-1):\n",
    "        below=1 \n",
    "print(p)\n",
    "print(a)\n",
    "print(max_len)\n",
    "\n",
    "\n",
    "X = []\n",
    "Y = []\n",
    "a = Fraction(a)\n",
    "for i in range(500000):\n",
    "    x = Fraction(random.randint(1, p-1))\n",
    "    result = bin(int((a*x)%p))[2:]\n",
    "    result = \"0\"*(max_len-len(result))+result\n",
    "    x_result = bin(int(x))[2:]\n",
    "    x_result = \"0\"*(max_len-len(x_result))+x_result\n",
    "    X.append([int(el) for el in result])\n",
    "    Y.append([int(el) for el in x_result])\n",
    "    \n",
    "X = np.array(X, dtype='float32')\n",
    "Y = np.array(Y, dtype='float32')\n",
    "print(X.shape)\n",
    "print(Y.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "134985a3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from sklearn.model_selection import train_test_split\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "\n",
    "class BaseModel1(nn.Module):\n",
    "    def __init__(self, ):\n",
    "        super(BaseModel1, self).__init__()\n",
    "        self.linear1 = nn.Linear(max_len, 1000) \n",
    "        self.linear2 = nn.Linear(1000, 1000)\n",
    "        self.linear3 = nn.Linear(1000, max_len)\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        h1 = torch.sigmoid(self.linear1(inputs))\n",
    "        h2 = torch.sigmoid(self.linear2(h1))\n",
    "        logits = self.linear3(h2)\n",
    "        return logits \n",
    "\n",
    "batch_size = 1000\n",
    "\n",
    "X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.3, random_state=42)\n",
    "\n",
    "X_train = torch.from_numpy(X_train)\n",
    "y_train = torch.from_numpy(y_train)\n",
    "X_test = torch.from_numpy(X_test) \n",
    "y_test = torch.from_numpy(y_test)\n",
    "\n",
    "\n",
    "train_dataset = TensorDataset(X_train,y_train)\n",
    "train_loader = DataLoader(train_dataset, batch_size = batch_size)\n",
    "\n",
    "test_dataset = TensorDataset(X_test,y_test)\n",
    "test_loader = DataLoader(test_dataset, batch_size = batch_size)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99b56649-7723-4972-abcc-7ba6e07bcd3e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "EPOCHS = 6000\n",
    "model = BaseModel1()\n",
    "model.cuda()\n",
    "criterion = torch.nn.BCEWithLogitsLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "data = []\n",
    "train_chunks = int(np.ceil(len(X_train)/batch_size))\n",
    "test_chunks = int(np.ceil(len(X_test)/batch_size))\n",
    "\n",
    "#train model\n",
    "for i in range(EPOCHS):\n",
    "    train_correct = 0\n",
    "    total_train_loss = 0\n",
    "    total_test_loss = 0\n",
    "    test_correct = 0\n",
    "    train_bit_correct = np.zeros(max_len)\n",
    "    test_bit_correct = np.zeros(max_len)\n",
    "    for j in range(train_chunks):\n",
    "        x = X_train[j*batch_size:(j+1)*batch_size]\n",
    "        y = y_train[j*batch_size:(j+1)*batch_size]\n",
    "\n",
    "        model.zero_grad()\n",
    "        optimizer.zero_grad()\n",
    "        x = x.cuda()\n",
    "        y = y.cuda()    \n",
    "        logits = model(x)\n",
    "        #sum of binary cross-entropy losses\n",
    "        loss = criterion(logits[:,0], y[:,0])\n",
    "        for z in range(1, max_len):\n",
    "            loss += criterion(logits[:,z], y[:,z])\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_train_loss += loss.item()\n",
    "        train_correct += ((y==(logits>0).to(float)).sum(axis=1)/max_len).sum().item()\n",
    "        train_bit_correct += (y==(logits>0).to(float)).sum(axis=0).cpu().numpy()\n",
    "        \n",
    "\n",
    "    with torch.no_grad():\n",
    "        for j in range(test_chunks):\n",
    "            x = X_test[j*batch_size:(j+1)*batch_size]\n",
    "            y = y_test[j*batch_size:(j+1)*batch_size]\n",
    "            x = x.cuda()\n",
    "            y = y.cuda()    \n",
    "\n",
    "            logits = model(x)\n",
    "            loss = criterion(logits[:,0], y[:,0])\n",
    "            for z in range(1, max_len):\n",
    "                loss += criterion(logits[:,z], y[:,z])\n",
    "            total_test_loss += loss.item()\n",
    "            test_correct += ((y==(logits>0).to(float)).sum(axis=1)/max_len).sum().item()\n",
    "            test_bit_correct += (y==(logits>0).to(float)).sum(axis=0).cpu().numpy()\n",
    "\n",
    "\n",
    "    train_acc = train_correct/len(X_train)\n",
    "    train_loss = total_train_loss/len(X_train)\n",
    "    test_loss = total_test_loss/len(X_test)\n",
    "    test_acc = test_correct/len(X_test)\n",
    "    data.append([train_acc, train_loss, test_loss, test_acc]+list(train_bit_correct/len(X_train))+list(test_bit_correct/len(X_test)))\n",
    "    if i % 20==0:\n",
    "        print(i,\"train_acc:\", round(train_acc, 6), 'train_loss:',\n",
    "              round(train_loss, 7), 'val_loss:', round(test_loss, 6), \n",
    "              \"val_acc:\", round(test_acc, 6))\n",
    "    if (test_acc>=1) and (train_acc>=1):\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa988886-ed30-48e2-b9ff-4ba93ae62cbd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "bit_train_cols = ['bit_train_'+str(i) for i in list(range(max_len,0,-1))]\n",
    "bit_test_cols = ['bit '+str(i) if i>=36 else \"_bit \"+str(i) for i in list(range(max_len,0,-1))]\n",
    "df = pd.DataFrame(data, columns = ['train_acc', 'train_loss', 'test_loss', 'test_acc']\\\n",
    "                 + bit_train_cols + bit_test_cols)\n",
    "df = df.reset_index().rename({'index':'step'}, axis=1)\n",
    "\n",
    "ax = df.plot(x='step', y=bit_test_cols)\n",
    "plt.xlabel('Step')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.rcParams.update({'font.size': 14})\n",
    "ax.legend(bbox_to_anchor=(1.22, 1.02), prop = { \"size\": 10 })\n",
    "plt.savefig('40_bits.png', dpi=300, bbox_inches = 'tight')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "KNNLM",
   "language": "python",
   "name": "knnlm"
  },
  "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}