Implemented GPU support.

src/lambo/examples/run_pretraining.py

@@ -6,6 +6,7 @@ import importlib_resources as resources
 from pathlib import Path
 
 import torch
+import sys
 
 from lambo.learning.dictionary import create_dictionary
 from lambo.learning.model_pretraining import LamboPretrainingNetwork
@@ -15,11 +16,13 @@ from lambo.learning.train import pretrain
 from lambo.utils.oscar import read_jsonl_to_documents, download_archive1_from_oscar
 
 if __name__=='__main__':
-    outpath = Path.home() / 'PATH-TO/models/pretrained/'
-    tmppath = Path.home() / 'PATH-TO/tmp/tmp.jsonl.gz'
+    outpath = sys.argv[1] #Path.home() / 'PATH-TO/models/pretrained/'
+    tmppath = sys.argv[2] #Path.home() / 'PATH-TO/tmp/tmp.jsonl.gz'
     # These need to be filled ine before running. OSCAR is avaialable on request.
-    OSCAR_LOGIN = ''
-    OSCAR_PASSWORD = ''
+    OSCAR_LOGIN = sys.argv[3]
+    OSCAR_PASSWORD = sys.argv[4]
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     
     languages_file_str = resources.read_text('lambo.resources', 'languages.txt', encoding='utf-8', errors='strict')
     languages = [line.split(' ')[1] for line in languages_file_str.split('\n') if
@@ -50,7 +53,7 @@ if __name__=='__main__':
             Xchars, Xutfs, Xmasks, Yvecs = encode_pretraining([document_train], dict, CONTEXT_LEN)
             _, train_dataloader, test_dataloader = prepare_dataloaders_pretraining([document_train],
                                                                                    [document_test], CONTEXT_LEN, 32, dict)
-            pretrain(model, train_dataloader, test_dataloader, 1)
+            pretrain(model, train_dataloader, test_dataloader, 1, device)
         torch.save(model, outpath / ('oscar_' + language + '.pth'))
         with open(outpath / ('oscar_' + language + '.dict'), "w") as file1:
             file1.writelines([x + '\t' + str(dict[x]) + '\n' for x in dict])

src/lambo/learning/train.py

@@ -9,17 +9,19 @@ from lambo.learning.preprocessing_dict import utf_category_dictionary, prepare_d
 from lambo.utils.ud_reader import read_treebank
 
 
-def train_loop(dataloader, model, optimizer):
+def train_loop(dataloader, model, optimizer, device='cpu'):
     """
     Training loop.
     
     :param dataloader: dataloader with training data
     :param model: model to be optimised
     :param optimizer: optimiser used
+    :param device: the device to use for computation
     :return: no value returned
     """
     size = len(dataloader.dataset)
     for batch, XY in enumerate(dataloader):
+        XY = [xy.to(device) for xy in XY]
         Xs = XY[:-1]
         Y = XY[-1]
         pred = model(*Xs)
@@ -32,12 +34,13 @@ def train_loop(dataloader, model, optimizer):
             print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
 
 
-def test_loop(dataloader, model):
+def test_loop(dataloader, model, device='cpu'):
     """
     Test loop.
     
     :param dataloader: dataloader with test data
     :param model: model to be tested
+    :param device: the device to use for computation
     :return: no value returned
     """
     num_batches = len(dataloader)
@@ -46,6 +49,7 @@ def test_loop(dataloader, model):
     size = [0, 0, 0, 0]
     with torch.no_grad():
         for XY in dataloader:
+            XY = [xy.to(device) for xy in XY]
             Xs = XY[:-1]
             Y = XY[-1]
             pred = model(*Xs)
@@ -63,12 +67,13 @@ def test_loop(dataloader, model):
         f"Test Error: \n Accuracy chars: {(100 * (correct[0] / size[0])):>5f}%, tokens: {(100 * (correct[1] / size[1])):>5f}%, mwtokens: {(100 * (correct[2] / size[2])):>5f}%, sentences: {(100 * (correct[3] / size[3])):>5f}%, Avg loss: {test_loss:>8f} \n")
 
 
-def test_loop_pretraining(dataloader, model):
+def test_loop_pretraining(dataloader, model, device='cpu'):
     """
     Test loop for pretraining.
 
     :param dataloader: dataloader with test data
     :param model: model to be tested
+    :param device: the device to use for computation
     :return: no value returned
     """
     num_batches = len(dataloader)
@@ -77,6 +82,7 @@ def test_loop_pretraining(dataloader, model):
     size = [0, 0]
     with torch.no_grad():
         for XY in dataloader:
+            XY = [xy.to(device) for xy in XY]
             Xs = XY[:-1]
             Y = XY[-1]
             pred = model(*Xs)
@@ -93,7 +99,7 @@ def test_loop_pretraining(dataloader, model):
         f"Test Error: \n Accuracy nontrivial: {(100 * (correct[0] / size[0])):>5f}%, trivial: {(100 * (correct[1] / size[1])):>5f}%, Avg loss: {test_loss:>8f} \n")
 
 
-def train_new_and_save(model_name, treebank_path, save_path, epochs=10):
+def train_new_and_save(model_name, treebank_path, save_path, epochs=10, device='cpu'):
     """
     Train a new LAMBO model and save it in filesystem.
     
@@ -101,6 +107,7 @@ def train_new_and_save(model_name, treebank_path, save_path, epochs=10):
     :param treebank_path: path to the treebank training data
     :param save_path: path to save the generated model
     :param epochs: number of epochs to run for (default: 10)
+    :param device: the device to use for computation
     :return: no value returned
     """
     if model_name not in ['LAMBO-BILSTM']:
@@ -120,7 +127,7 @@ def train_new_and_save(model_name, treebank_path, save_path, epochs=10):
                                                                            BATCH_SIZE)
     model = LamboNetwork(MAX_LEN, dict, len(utf_category_dictionary))
     
-    tune(model, train_dataloader, test_dataloader, epochs)
+    tune(model, train_dataloader, test_dataloader, epochs, device)
     
     print("Saving")
     torch.save(model, save_path / (treebank_path.name + '.pth'))
@@ -128,7 +135,7 @@ def train_new_and_save(model_name, treebank_path, save_path, epochs=10):
         file1.writelines([x + '\t' + str(dict[x]) + '\n' for x in dict])
 
 
-def train_pretrained_and_save(language, treebank_path, save_path, pretrained_path, epochs=10):
+def train_pretrained_and_save(language, treebank_path, save_path, pretrained_path, epochs=10, device='cpu'):
     """
     Train a new LAMBO model, staring from pretrained, and save it in filesystem.
 
@@ -137,6 +144,7 @@ def train_pretrained_and_save(language, treebank_path, save_path, pretrained_pat
     :param save_path: path to save the generated model
     :param pretrained_path: path to the pretraining models
     :param epochs: number of epochs to run for (default: 10)
+    :param device: the device to use for computation
     :return: no value returned
     """
     print("Loading pretrained model")
@@ -166,7 +174,7 @@ def train_pretrained_and_save(language, treebank_path, save_path, pretrained_pat
                                                                            MAX_LEN,
                                                                            BATCH_SIZE, dict=dict)
     
-    tune(model, train_dataloader, test_dataloader, epochs)
+    tune(model, train_dataloader, test_dataloader, epochs, device)
     
     print("Saving")
     torch.save(model, save_path / (treebank_path.name + '.pth'))
@@ -174,7 +182,7 @@ def train_pretrained_and_save(language, treebank_path, save_path, pretrained_pat
         file1.writelines([x + '\t' + str(dict[x]) + '\n' for x in dict])
 
 
-def tune(model, train_dataloader, test_dataloader, epochs):
+def tune(model, train_dataloader, test_dataloader, epochs, device='cpu'):
     """
     Tune an existing LAMBO model with the provided data
     
@@ -182,9 +190,11 @@ def tune(model, train_dataloader, test_dataloader, epochs):
     :param train_dataloader: dataloader for training data
     :param test_dataloader: dataloader for test data
     :param epochs: number of epochs to run for
+    :param device: the device to use for computation
     :return: no value returned
     """
     print("Preparing training")
+    model.to(device)
     learning_rate = 1e-3
     optimizer = Adam(model.parameters(), lr=learning_rate)
     
@@ -192,11 +202,11 @@ def tune(model, train_dataloader, test_dataloader, epochs):
     test_loop(test_dataloader, model)
     for t in range(epochs):
         print(f"Epoch {t + 1}\n-------------------------------")
-        train_loop(train_dataloader, model, optimizer)
-        test_loop(test_dataloader, model)
+        train_loop(train_dataloader, model, optimizer, device)
+        test_loop(test_dataloader, model, device)
 
 
-def pretrain(model, train_dataloader, test_dataloader, epochs):
+def pretrain(model, train_dataloader, test_dataloader, epochs, device='cpu'):
     """
     Tune an existing LAMBO pretraining model with the provided data
     
@@ -204,9 +214,11 @@ def pretrain(model, train_dataloader, test_dataloader, epochs):
     :param train_dataloader: dataloader for training data
     :param test_dataloader: dataloader for test data
     :param epochs: number of epochs to run for
+    :param device: the device to use for computation
     :return: no value returned
     """
     print("Preparing pretraining")
+    model.to(device)
     learning_rate = 1e-3
     optimizer = Adam(model.parameters(), lr=learning_rate)
     
@@ -214,5 +226,5 @@ def pretrain(model, train_dataloader, test_dataloader, epochs):
     test_loop_pretraining(test_dataloader, model)
     for t in range(epochs):
         print(f"Epoch {t + 1}\n-------------------------------")
-        train_loop(train_dataloader, model, optimizer)
-        test_loop_pretraining(test_dataloader, model)
+        train_loop(train_dataloader, model, optimizer, device)
+        test_loop_pretraining(test_dataloader, model, device)