Support specifying which GPU

This commit is contained in:
Matthew Honnibal 2017-06-03 16:10:23 -05:00
parent 3f5c85d8de
commit 21eef90dbc
2 changed files with 10 additions and 4 deletions

View File

@ -27,14 +27,14 @@ from .. import displacy
dev_data=("location of JSON-formatted development data (optional)", "positional", None, str), dev_data=("location of JSON-formatted development data (optional)", "positional", None, str),
n_iter=("number of iterations", "option", "n", int), n_iter=("number of iterations", "option", "n", int),
n_sents=("number of sentences", "option", "ns", int), n_sents=("number of sentences", "option", "ns", int),
use_gpu=("Use GPU", "flag", "G", bool), use_gpu=("Use GPU", "option", "g", int),
resume=("Whether to resume training", "flag", "R", bool), resume=("Whether to resume training", "flag", "R", bool),
no_tagger=("Don't train tagger", "flag", "T", bool), no_tagger=("Don't train tagger", "flag", "T", bool),
no_parser=("Don't train parser", "flag", "P", bool), no_parser=("Don't train parser", "flag", "P", bool),
no_entities=("Don't train NER", "flag", "N", bool) no_entities=("Don't train NER", "flag", "N", bool)
) )
def train(cmd, lang, output_dir, train_data, dev_data, n_iter=20, n_sents=0, def train(cmd, lang, output_dir, train_data, dev_data, n_iter=20, n_sents=0,
use_gpu=False, resume=False, no_tagger=False, no_parser=False, no_entities=False): use_gpu=-1, resume=False, no_tagger=False, no_parser=False, no_entities=False):
""" """
Train a model. Expects data in spaCy's JSON format. Train a model. Expects data in spaCy's JSON format.
""" """
@ -76,7 +76,7 @@ def train(cmd, lang, output_dir, train_data, dev_data, n_iter=20, n_sents=0,
corpus = GoldCorpus(train_path, dev_path, limit=n_sents) corpus = GoldCorpus(train_path, dev_path, limit=n_sents)
n_train_docs = corpus.count_train() n_train_docs = corpus.count_train()
optimizer = nlp.begin_training(lambda: corpus.train_tuples, use_gpu=use_gpu) optimizer = nlp.begin_training(lambda: corpus.train_tuples, device=use_gpu)
print("Itn.\tLoss\tUAS\tNER P.\tNER R.\tNER F.\tTag %\tToken %") print("Itn.\tLoss\tUAS\tNER P.\tNER R.\tNER F.\tTag %\tToken %")
try: try:

View File

@ -279,9 +279,14 @@ class Language(object):
for word in annots[1]: for word in annots[1]:
_ = self.vocab[word] _ = self.vocab[word]
contexts = [] contexts = []
if cfg.get('use_gpu'): if cfg.get('device', -1) >= 0:
import cupy.cuda.device
device = cupy.cuda.device.Device(cfg['device'])
device.use()
Model.ops = CupyOps() Model.ops = CupyOps()
Model.Ops = CupyOps Model.Ops = CupyOps
else:
device = None
for proc in self.pipeline: for proc in self.pipeline:
if hasattr(proc, 'begin_training'): if hasattr(proc, 'begin_training'):
context = proc.begin_training(get_gold_tuples(), context = proc.begin_training(get_gold_tuples(),
@ -296,6 +301,7 @@ class Language(object):
optimizer = Adam(Model.ops, learn_rate, L2=L2, beta1=beta1, optimizer = Adam(Model.ops, learn_rate, L2=L2, beta1=beta1,
beta2=beta2, eps=eps) beta2=beta2, eps=eps)
optimizer.max_grad_norm = max_grad_norm optimizer.max_grad_norm = max_grad_norm
optimizer.device = device
return optimizer return optimizer
def evaluate(self, docs_golds): def evaluate(self, docs_golds):