Skip to content
Snippets Groups Projects
Commit 1986518e authored by piotrmp's avatar piotrmp
Browse files

Implemented GPU support.

parent 8e0d8500
No related branches found
No related tags found
1 merge request!1Migration to UD 2.11
Show whitespace changes
Inline Side-by-side
src/lambo/examples/run_pretraining.py
+ 8
− 5
View file @ 1986518e
......@@ -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
+ 25
− 13
View file @ 1986518e
......@@ -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)
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment