Partial GPU support in images

docs/registry.json

@@ -121,9 +121,9 @@
     "ref_url": "https://github.com/yikangshen/Ordered-Neurons",
     "image": {
       "maintainer": "[email protected]",
-      "version": "ca867981a9bd1ac987447ab6f77e0850236f604d",
-      "checksum": "f302550ef5bd6f45c6adbbd98eefc56db7f0dd73",
-      "datetime": "2020-05-25T17:43:02.573233114Z",
+      "version": "9d8db666f1fe9fa22cbd37ab6eff7c15a0fb6a8a",
+      "checksum": "8e99c084ea20275e9a247badf99139a8a20eb6c7",
+      "datetime": "2020-05-26T18:39:36.015922392Z",
       "supported_features": {
         "tokenize": true,
         "unkify": true,
@@ -133,12 +133,12 @@
       },
       "gpu": {
         "required": false,
-        "supported": false
+        "supported": true
       },
       "registry": "docker.io",
       "name": "cpllab/language-models",
       "tag": "ordered-neurons",
-      "size": 3218739014
+      "size": 3218738720
     },
     "tokenizer": {
       "type": "word",
@@ -179,9 +179,9 @@
     "ref_url": "https://github.com/cpllab/tinylstm",
     "image": {
       "maintainer": "[email protected]",
-      "version": "75a97862cbbade0574d9a03ddd60103b7fe07363",
-      "checksum": "ac57eb3c8b19985978b61df9415e3a95be6ae7b4",
-      "datetime": "2020-05-26T18:32:16.22094886Z",
+      "version": "9d8db666f1fe9fa22cbd37ab6eff7c15a0fb6a8a",
+      "checksum": "344a20e0e67a8ee39c38e34ae1a015f2eb8e99d3",
+      "datetime": "2020-05-26T18:39:22.988852042Z",
       "supported_features": {
         "tokenize": true,
         "unkify": true,
@@ -191,12 +191,12 @@
       },
       "gpu": {
         "required": false,
-        "supported": false
+        "supported": true
       },
       "registry": "docker.io",
       "name": "cpllab/language-models",
       "tag": "tinylstm",
-      "size": 3295284385
+      "size": 3295284357
     },
     "tokenizer": {
       "type": "word",

models/ordered-neurons/bin/get_surprisals

@@ -8,4 +8,4 @@ INPUT_FILE="$1"
 
 python ${MODEL_ROOT}/get_surprisals.py \
     "$MODEL_CHECKPOINT" /tmp/input_tokenized \
-    --corpus_file "$MODEL_CORPUS" --no-cuda
+    --corpus_file "$MODEL_CORPUS"

models/ordered-neurons/get_surprisals.py

@@ -1,4 +1,5 @@
 import argparse
+import os
 from pathlib import Path
 import pickle
 import sys
@@ -20,30 +21,26 @@
 parser.add_argument("--bptt", type=int, default=70, help="sequence length")
 parser.add_argument("--emsize", type=int, default=400, help="size of word embeddings")
 parser.add_argument("--seed", type=int, default=1111, help="random seed")
-parser.add_argument("--cuda", dest="cuda", action="store_true")
-parser.add_argument("--no-cuda", dest="cuda", action="store_false")
-parser.set_defaults(cuda=True)
 
 args = parser.parse_args()
 
+use_cuda = torch.cuda.is_available() and os.environ.get("LMZOO_USE_GPU", False)
+if use_cuda:
+    sys.stderr.write("Using GPU device.\n")
+device = "cuda" if use_cuda else "cpu"
+
 # Set the random seed manually for reproducibility.
 def set_seed(seed):
     np.random.seed(seed)
     torch.manual_seed(seed)
-    if torch.cuda.is_available():
-        if not args.cuda:
-            print("WARNING: You have a CUDA device, so you should probably run with --cuda")
-        else:
-            torch.cuda.manual_seed(seed)
+    if use_cuda:
+        torch.cuda.manual_seed(seed)
 
 
 def model_load(fn):
     global model, criterion, optimizer
     with open(fn, 'rb') as f:
-        kwargs = {}
-        if not args.cuda:
-            kwargs["map_location"] = "cpu"
-        model, criterion, optimizer = torch.load(f, **kwargs)
+        model, criterion, optimizer = torch.load(f, map_location=device)
 
 
 def get_batch(data_source, i, window):
@@ -73,7 +70,7 @@ def get_surprisals(sentences, corpus, outf, seed):
                 raise RuntimeError("Internal error: Dictionary lookup failed. This should not happen with properly unked inputs.")
 
         # model expects T * batch_size array
-        data_source = data_source.unsqueeze(1)
+        data_source = data_source.unsqueeze(1).to(device)
 
         with torch.no_grad():
             hidden = model.init_hidden(1)
@@ -84,19 +81,15 @@ def get_surprisals(sentences, corpus, outf, seed):
                         torch.nn.functional.linear(output, model.decoder.weight, bias=model.decoder.bias),
                         dim=1)
 
-                # Convert to numpy and change to log2.
-                logprobs = logprobs.detach().numpy()
-                logprobs /= np.log(2)
-
-                # Retrieve relevant surprisal values.
-                targets = targets.numpy()
-                target_surprisals = -logprobs[np.arange(len(targets)), targets]
+                # Convert to surprisals and extract relevant surprisal value.
+                surprisals = - logprobs / np.log(2)
+                target_surprisals = surprisals[np.arange(len(targets)), targets].cpu()
 
                 for k, surp in enumerate(target_surprisals):
                     outf.write("%i\t%i\t%s\t%f\n" % (i + 1, j + k + 2, sentence[j + k + 1], surp))
 
 
-corpus = torch.load(args.corpus_file)
+corpus = torch.load(args.corpus_file, map_location=device)
 model_load(args.model_checkpoint)
 sentences = [line.strip().split(" ") for line in args.file.readlines() if line.strip()]
 get_surprisals(sentences, corpus, args.outf, args.seed)

models/ordered-neurons/spec.template.json

@@ -16,7 +16,7 @@
         },
         "gpu": {
             "required": false,
-            "supported": false
+            "supported": true
         }
     },
 

models/tinylstm/Dockerfile

@@ -31,6 +31,7 @@ ARG MODEL_ROOT=models/tinylstm
 COPY --from=builder /opt/tinylstm /opt/tinylstm
 
 RUN pip install -q nltk
+COPY ${MODEL_ROOT}/get_surprisals.py /opt/tinylstm
 
 # Add test dependencies.
 RUN pip install nose jsonschema

models/tinylstm/bin/get_surprisals

@@ -6,7 +6,7 @@ INPUT_FILE="$1"
 
 /opt/bin/tokenize $1 > /tmp/input_tokenized
 
-python ${MODEL_ROOT}/eval.py \
+python ${MODEL_ROOT}/get_surprisals.py \
     --checkpoint "$MODEL_CHECKPOINT" \
     --corpus "$MODEL_CORPUS" \
     --eval_data /tmp/input_tokenized \

models/tinylstm/get_surprisals.py

@@ -0,0 +1,86 @@
+import argparse
+import os
+import sys
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+sys.path.append("/opt/tinylstm")
+import data
+
+parser = argparse.ArgumentParser(description='PyTorch PTB Language Model')
+
+# Model parameters.
+parser.add_argument('--corpus', type=str, default=None, required=True,
+                    help='location of corpus file')
+parser.add_argument('--checkpoint', type=str, required=True,
+                    help='model checkpoint to use')
+parser.add_argument('--seed', type=int, default=1111,
+                    help='random seed')
+parser.add_argument('--eval_data', type=str, default='stimuli_items/input_test.raw')
+parser.add_argument('--outf', type=argparse.FileType("w", encoding="utf-8"), default=sys.stdout,
+                    help='output file for generated text')
+
+parser.set_defaults(refresh_state=True)
+parser.add_argument("--no_refresh_state", dest="refresh_state", action="store_false",
+                    help="Don't refresh the RNN hidden state between sentences.")
+
+args = parser.parse_args()
+
+use_cuda = torch.cuda.is_available() and bool(os.environ.get("LMZOO_USE_GPU", False))
+if use_cuda:
+    sys.stderr.write("Using GPU device.\n")
+
+# Set the random seed manually for reproducibility.
+torch.manual_seed(args.seed)
+if use_cuda:
+    torch.cuda.manual_seed(args.seed)
+
+device = torch.device("cuda" if use_cuda else "cpu")
+
+with open(args.checkpoint, 'rb') as f:
+    if use_cuda:
+        model = torch.load(f).to(device)
+    else:
+        model = torch.load(f, map_location=lambda storage, loc: storage)
+        model.cpu()
+model.eval()
+
+
+corpus = torch.load(args.corpus)
+
+ntokens = len(corpus.dictionary)
+hidden = model.init_hidden(1)
+input = torch.randint(ntokens, (1, 1), dtype=torch.long).to(device)
+
+
+# read eval data
+with open(args.eval_data, 'r') as f:
+    lines = f.readlines()
+sents = [line.strip().split() for line in lines]
+
+
+with args.outf as f:
+    # write header.
+    f.write("sentence_id\ttoken_id\ttoken\tsurprisal\n")
+    with torch.no_grad():  # no tracking history
+        # all_ppls = []
+        for sent_id, sent in enumerate(tqdm(sents)):
+            if args.refresh_state:
+                hidden = model.init_hidden(1)
+
+            input = torch.tensor([[corpus.dictionary.word2idx[sent[0]]]],dtype=torch.long).to(device)
+
+            f.write("%i\t%i\t%s\t%f\n" % (sent_id + 1, 1, sent[0], 0.0))
+
+            for i, w in enumerate(sent[1:]):
+                output, hidden = model(input, hidden)
+                surprisals = - F.log_softmax(output, dim=2) / np.log(2)
+                word_idx = corpus.dictionary.word2idx[w]
+                word_surprisal = surprisals[0, 0, word_idx].item()
+
+                f.write("%i\t%i\t%s\t%f\n" % (sent_id + 1, i + 2, w, word_surprisal))
+
+                input.fill_(word_idx)

models/tinylstm/spec.template.json

@@ -16,7 +16,7 @@
         },
         "gpu": {
             "required": false,
-            "supported": false
+            "supported": true
         }
     },