Make character copy non-blocking

This commit is contained in:
Matthw Honnibal 2019-10-18 17:23:19 +02:00
parent 4da1c1c211
commit 727ede6599

View File

@ -984,15 +984,23 @@ class CharacterEmbed(Model):
def begin_update(self, docs, drop=0.0): def begin_update(self, docs, drop=0.0):
if not docs: if not docs:
return [] return []
ids = []
output = [] output = []
weights = self.vectors weights = self.vectors
# This assists in indexing; it's like looping over this dimension. # This assists in indexing; it's like looping over this dimension.
# Still consider this weird witch craft...But thanks to Mark Neumann # Still consider this weird witch craft...But thanks to Mark Neumann
# for the tip. # for the tip.
nCv = self.ops.xp.arange(self.nC) nCv = self.ops.xp.arange(self.nC)
stream = util.get_cuda_stream(non_blocking=True)
ids = []
for doc in docs: for doc in docs:
doc_ids = self.ops.asarray(doc.to_utf8_array(nr_char=self.nC)) doc_ids = doc.to_utf8_array(nr_char=self.nC)
if stream is None:
ids.append(doc_ids)
else:
ids.append(util.get_async(stream, doc_ids))
if stream is not None:
stream.synchronize()
for doc, doc_ids in zip(docs, ids):
doc_vectors = self.ops.allocate((len(doc), self.nC, self.nM)) doc_vectors = self.ops.allocate((len(doc), self.nC, self.nM))
# Let's say I have a 2d array of indices, and a 3d table of data. What numpy # Let's say I have a 2d array of indices, and a 3d table of data. What numpy
# incantation do I chant to get # incantation do I chant to get