explosion/spaCy
1
1
Fork 0
mirror of https://github.com/explosion/spaCy.git synced 2024-09-21 03:19:13 +03:00
Code Issues Packages Projects Releases Wiki Activity

Add util function to enable GPU

Browse Source
This commit is contained in:
Matthew Honnibal 2017-09-20 19:16:35 -05:00
parent 24e85c2048
commit ffda38356a
1 changed files with 12 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
spacy/util.py
View File

@ -14,6 +14,7 @@ import numpy
import io
import dill
from collections import OrderedDict
from thinc.neural._classes.model import Model
import msgpack
import msgpack_numpy
@ -557,3 +558,14 @@ def minify_html(html):
RETURNS (unicode): "Minified" HTML.
"""
return html.strip().replace(' ', '').replace('\n', '')
def use_gpu(gpu_id):
import cupy.cuda.device
from thinc.neural.ops import CupyOps
device = cupy.cuda.device.Device(gpu_id)
device.use()
Model.ops = CupyOps()
Model.Ops = CupyOps
return device