Rooba
/
spaCy
mirror of https://github.com/explosion/spaCy.git
1
1
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Set GPU before loading any models in train CLI (#4989) 

Set the GPU before loading any existing models in the train CLI so that
you can start with a base model and train on GPU.
This commit is contained in:
adrianeboyd 2020-02-11 23:45:41 +01:00 committed by GitHub
parent 842dfddbb9
commit 99a543367d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 13 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
spacy/cli/train.py
View File

@ -14,6 +14,7 @@ import contextlib
import random
from .._ml import create_default_optimizer
from ..util import use_gpu as set_gpu
from ..attrs import PROB, IS_OOV, CLUSTER, LANG
from ..gold import GoldCorpus
from ..compat import path2str
@ -147,6 +148,18 @@ def train(
disabled_pipes = None
pipes_added = False
msg.text("Training pipeline: {}".format(pipeline))
if use_gpu >= 0:
activated_gpu = None
try:
activated_gpu = set_gpu(use_gpu)
except Exception as e:
msg.warn("Exception: {}".format(e))
if activated_gpu is not None:
msg.text("Using GPU: {}".format(use_gpu))
else:
msg.warn("Unable to activate GPU: {}".format(use_gpu))
msg.text("Using CPU only")
use_gpu = -1
if base_model:
msg.text("Starting with base model '{}'".format(base_model))
nlp = util.load_model(base_model)