smaug_train_wholegene.ipynb

{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"name":"smaug_train_wholegene.ipynb","provenance":[],"collapsed_sections":[]},"kernelspec":{"name":"python3","display_name":"Python 3"},"accelerator":"GPU"},"cells":[{"cell_type":"code","metadata":{"id":"Ga8mM1hmmisE","colab_type":"code","outputId":"7427dd4e-3488-4d1c-b066-22111fe52b61","executionInfo":{"status":"ok","timestamp":1575583476643,"user_tz":300,"elapsed":2916,"user":{"displayName":"Markus Sommer","photoUrl":"","userId":"08083804319459268544"}},"colab":{"base_uri":"https://localhost:8080/","height":357}},"source":["!pip install transformers\n","# !pip install torch\n","# !pip install conda\n","# !conda install pytorch=0.4.1 cuda90 -c pytorch"],"execution_count":1,"outputs":[{"output_type":"stream","text":["Requirement already satisfied: transformers in /usr/local/lib/python3.6/dist-packages (2.2.1)\n","Requirement already satisfied: tqdm in /usr/local/lib/python3.6/dist-packages (from transformers) (4.28.1)\n","Requirement already satisfied: requests in /usr/local/lib/python3.6/dist-packages (from transformers) (2.21.0)\n","Requirement already satisfied: sacremoses in /usr/local/lib/python3.6/dist-packages (from transformers) (0.0.35)\n","Requirement already satisfied: sentencepiece in /usr/local/lib/python3.6/dist-packages (from transformers) (0.1.83)\n","Requirement already satisfied: regex in /usr/local/lib/python3.6/dist-packages (from transformers) (2019.11.1)\n","Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from transformers) (1.17.4)\n","Requirement already satisfied: boto3 in /usr/local/lib/python3.6/dist-packages (from transformers) (1.10.27)\n","Requirement already satisfied: idna<2.9,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (2.8)\n","Requirement already satisfied: chardet<3.1.0,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (3.0.4)\n","Requirement already satisfied: urllib3<1.25,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (1.24.3)\n","Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (2019.9.11)\n","Requirement already satisfied: joblib in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers) (0.14.0)\n","Requirement already satisfied: click in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers) (7.0)\n","Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers) (1.12.0)\n","Requirement already satisfied: s3transfer<0.3.0,>=0.2.0 in /usr/local/lib/python3.6/dist-packages (from boto3->transformers) (0.2.1)\n","Requirement already satisfied: jmespath<1.0.0,>=0.7.1 in /usr/local/lib/python3.6/dist-packages (from boto3->transformers) (0.9.4)\n","Requirement already satisfied: botocore<1.14.0,>=1.13.27 in /usr/local/lib/python3.6/dist-packages (from boto3->transformers) (1.13.27)\n","Requirement already satisfied: python-dateutil<2.8.1,>=2.1; python_version >= \"2.7\" in /usr/local/lib/python3.6/dist-packages (from botocore<1.14.0,>=1.13.27->boto3->transformers) (2.6.1)\n","Requirement already satisfied: docutils<0.16,>=0.10 in /usr/local/lib/python3.6/dist-packages (from botocore<1.14.0,>=1.13.27->boto3->transformers) (0.15.2)\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"LfDztkj3muhS","colab_type":"code","outputId":"bd34e418-da2b-4573-8702-542098521f64","executionInfo":{"status":"ok","timestamp":1575583476645,"user_tz":300,"elapsed":2866,"user":{"displayName":"Markus Sommer","photoUrl":"","userId":"08083804319459268544"}},"colab":{"base_uri":"https://localhost:8080/","height":34}},"source":["# Load the Drive helper and mount\n","from google.colab import drive\n","\n","# This will prompt for authorization.\n","drive.mount('/content/drive')"],"execution_count":2,"outputs":[{"output_type":"stream","text":["Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"nq_dBVgkmuer","colab_type":"code","outputId":"8eb94173-9fbd-4766-a3c5-5b01d98c07c4","executionInfo":{"status":"ok","timestamp":1575583478266,"user_tz":300,"elapsed":4462,"user":{"displayName":"Markus Sommer","photoUrl":"","userId":"08083804319459268544"}},"colab":{"base_uri":"https://localhost:8080/","height":63}},"source":["import os\n","import sys\n","import pickle\n","import torch\n","import numpy as np\n","import torch.optim as optim\n","# from transformers import BertConfig, BertModel, BertForMaskedLM\n","from transformers import GPT2Config, GPT2Model, GPT2LMHeadModel\n","\n","from IPython.display import clear_output"],"execution_count":3,"outputs":[{"output_type":"display_data","data":{"text/html":["<p style=\"color: red;\">\n","The default version of TensorFlow in Colab will soon switch to TensorFlow 2.x.<br>\n","We recommend you <a href=\"https://www.tensorflow.org/guide/migrate\" target=\"_blank\">upgrade</a> now \n","or ensure your notebook will continue to use TensorFlow 1.x via the <code>%tensorflow_version 1.x</code> magic:\n","<a href=\"https://colab.research.google.com/notebooks/tensorflow_version.ipynb\" target=\"_blank\">more info</a>.</p>\n"],"text/plain":["<IPython.core.display.HTML object>"]},"metadata":{"tags":[]}}]},{"cell_type":"code","metadata":{"id":"B9MlYOB5mudK","colab_type":"code","outputId":"b9cad896-ec9a-44bd-9601-d851b42df8c7","executionInfo":{"status":"ok","timestamp":1575583482135,"user_tz":300,"elapsed":8305,"user":{"displayName":"Markus Sommer","photoUrl":"","userId":"08083804319459268544"}},"colab":{"base_uri":"https://localhost:8080/","height":1000}},"source":["# custom parameters for GPT2 model\n","vocab_size = 21\n","max_position_embeddings = 1024 # 1024\n","n_ctw = max_position_embeddings # 1024\n","n_embd = int(768) # 768\n","n_layer = 6 # 12\n","n_head = 6 # 12\n","resid_pdrop = 0.0 # 0.1\n","embd_pdrop = 0.0 # 0.1\n","attn_pdrop = 0.0 # 0.1\n","layer_norm_epsilon = 1e-5 # 1e-5\n","\n","\n","config = GPT2Config(vocab_size_or_config_json_file=vocab_size,\n","                    n_positions=max_position_embeddings,\n","                    n_ctw=n_ctw,\n","                    n_embd=n_embd,\n","                    n_layer=n_layer,\n","                    n_head=n_head, \n","                    resid_pdrop=resid_pdrop,\n","                    embd_pdrop=embd_pdrop,\n","                    attn_pdrop=attn_pdrop,\n","                    layer_norm_epsilon=layer_norm_epsilon)\n","\n","model = GPT2LMHeadModel(config)\n","\n","# print(model)\n","model.to('cuda')"],"execution_count":4,"outputs":[{"output_type":"execute_result","data":{"text/plain":["GPT2LMHeadModel(\n","  (transformer): GPT2Model(\n","    (wte): Embedding(21, 768)\n","    (wpe): Embedding(1024, 768)\n","    (drop): Dropout(p=0.0, inplace=False)\n","    (h): ModuleList(\n","      (0): Block(\n","        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n","        (attn): Attention(\n","          (c_attn): Conv1D()\n","          (c_proj): Conv1D()\n","          (attn_dropout): Dropout(p=0.0, inplace=False)\n","          (resid_dropout): Dropout(p=0.0, inplace=False)\n","        )\n","        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n","        (mlp): MLP(\n","          (c_fc): Conv1D()\n","          (c_proj): Conv1D()\n","          (dropout): Dropout(p=0.0, inplace=False)\n","        )\n","      )\n","      (1): Block(\n","        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n","        (attn): Attention(\n","          (c_attn): Conv1D()\n","          (c_proj): Conv1D()\n","          (attn_dropout): Dropout(p=0.0, inplace=False)\n","          (resid_dropout): Dropout(p=0.0, inplace=False)\n","        )\n","        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n","        (mlp): MLP(\n","          (c_fc): Conv1D()\n","          (c_proj): Conv1D()\n","          (dropout): Dropout(p=0.0, inplace=False)\n","        )\n","      )\n","      (2): Block(\n","        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n","        (attn): Attention(\n","          (c_attn): Conv1D()\n","          (c_proj): Conv1D()\n","          (attn_dropout): Dropout(p=0.0, inplace=False)\n","          (resid_dropout): Dropout(p=0.0, inplace=False)\n","        )\n","        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n","        (mlp): MLP(\n","          (c_fc): Conv1D()\n","          (c_proj): Conv1D()\n","          (dropout): Dropout(p=0.0, inplace=False)\n","        )\n","      )\n","      (3): Block(\n","        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n","        (attn): Attention(\n","          (c_attn): Conv1D()\n","          (c_proj): Conv1D()\n","          (attn_dropout): Dropout(p=0.0, inplace=False)\n","          (resid_dropout): Dropout(p=0.0, inplace=False)\n","        )\n","        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n","        (mlp): MLP(\n","          (c_fc): Conv1D()\n","          (c_proj): Conv1D()\n","          (dropout): Dropout(p=0.0, inplace=False)\n","        )\n","      )\n","      (4): Block(\n","        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n","        (attn): Attention(\n","          (c_attn): Conv1D()\n","          (c_proj): Conv1D()\n","          (attn_dropout): Dropout(p=0.0, inplace=False)\n","          (resid_dropout): Dropout(p=0.0, inplace=False)\n","        )\n","        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n","        (mlp): MLP(\n","          (c_fc): Conv1D()\n","          (c_proj): Conv1D()\n","          (dropout): Dropout(p=0.0, inplace=False)\n","        )\n","      )\n","      (5): Block(\n","        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n","        (attn): Attention(\n","          (c_attn): Conv1D()\n","          (c_proj): Conv1D()\n","          (attn_dropout): Dropout(p=0.0, inplace=False)\n","          (resid_dropout): Dropout(p=0.0, inplace=False)\n","        )\n","        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n","        (mlp): MLP(\n","          (c_fc): Conv1D()\n","          (c_proj): Conv1D()\n","          (dropout): Dropout(p=0.0, inplace=False)\n","        )\n","      )\n","    )\n","    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n","  )\n","  (lm_head): Linear(in_features=768, out_features=21, bias=False)\n",")"]},"metadata":{"tags":[]},"execution_count":4}]},{"cell_type":"code","metadata":{"id":"l5Hy83Homuag","colab_type":"code","colab":{}},"source":["# load data\n","d = \"drive/My Drive/Colab Notebooks/smaug/data\"\n","shingle_path = os.path.join(d, \"ecoli_MG1655_shingles_length20_overlap10.npy\")\n","# shingle_path = os.path.join(d, \"ecoli_MG1655_1024aa_noshingle.npy\")\n","\n","\n","with open(shingle_path, 'rb') as f:\n","    ecoli_shingles = np.load(shingle_path)#[:2056]"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"id":"nPXHybK8muYV","colab_type":"code","colab":{}},"source":["# pass to GPU\n","dtrain = torch.tensor(ecoli_shingles).to('cuda')\n","#  = tokens_tensor"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"id":"I5r8Jf3pwpPX","colab_type":"code","colab":{}},"source":[""],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"id":"W4XFDO6wmuV8","colab_type":"code","colab":{}},"source":["lr = 1e-2 # 1e-3\n","optimizer = optim.AdamW(model.parameters(), lr=lr)"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"id":"T85Z_nICmuTq","colab_type":"code","colab":{}},"source":["def savemodel():\n","    # save model\n","    # Load the Drive helper and mount\n","    from google.colab import drive\n","\n","    # This will prompt for authorization.\n","    drive.mount('/content/drive')\n","\n","    modeldir = \"drive/My Drive/Colab Notebooks/smaug/data/models/ecoli_trivial_wholegene\"\n","    model.save_pretrained(modeldir)\n","    print(os.listdir(modeldir))"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"id":"lpCjshJ5muSY","colab_type":"code","outputId":"2d268f45-24d7-40b9-9b02-4800ad191918","colab":{"base_uri":"https://localhost:8080/","height":119}},"source":["# train model with single aa masked at a time ##### batch\n","model.train()\n","\n","batch_size = 2**8\n","\n","optimizer.zero_grad()\n","# np.random.seed(42424)\n","for i in range(1000000):\n","    if i%500==10:\n","        savemodel()\n","    optimizer.zero_grad()\n","\n","    select_idx = np.random.randint(0, len(dtrain), batch_size)\n","\n","#     input_ids = dtrain[select_idx].unsqueeze(0) # singleton\n","    input_ids = dtrain[select_idx]\n","\n","    attention_mask = torch.tensor((input_ids != 0), dtype=float).to('cuda') # ignore padding 0's\n","    \n","    outputs = model(input_ids, labels=input_ids, attention_mask=attention_mask)\n","    loss, prediction_scores = outputs[:2]\n","    \n","    loss.backward()\n","    optimizer.step()\n","    \n","    clear_output(wait=True)\n","    print(\"Loss:\", loss.item())\n","\n","    for k in range(0, 5):\n","#         print(i, torch.argmax(prediction_scores[0,k-1]).item(), input_ids[0,k].item(), \"\\t\", loss.item()) #TODO figure out why GPT2 only offsets sometimes\n","        print(i, \"\\t\", torch.argmax(prediction_scores[0,k-1]).item(), \"\\t\", input_ids[0,k].item(), \"\\t\")\n","#         print(prediction_scores[0])\n","    \n","    # break"],"execution_count":0,"outputs":[{"output_type":"stream","text":["Loss: 1.2043142318725586\n","78698 \t 17 \t 4 \t\n","78698 \t 1 \t 1 \t\n","78698 \t 8 \t 8 \t\n","78698 \t 1 \t 1 \t\n","78698 \t 7 \t 7 \t\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"tRm7QgtSmuLV","colab_type":"code","colab":{}},"source":[""],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"id":"NhYViNqFmuIx","colab_type":"code","colab":{}},"source":[""],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"id":"Cm_KUnUPmuFM","colab_type":"code","colab":{}},"source":[""],"execution_count":0,"outputs":[]}]}