explosion/spaCy
1
1
Fork 0
mirror of https://github.com/explosion/spaCy.git synced 2025-10-30 15:37:29 +03:00
Code Issues Packages Projects Releases Wiki Activity

Use scatter_add to speed up span embed backprop

Browse Source 
This was the slowest part of the code, and using scatter_add here
probably reduces the runtime by 50%.
This commit is contained in:
Paul O'Leary McCann 2021-07-10 18:08:51 +09:00
parent d0b041aff4
commit f34915c1e8
1 changed files with 4 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
spacy/ml/models/coref.py
View File

@ -187,14 +187,10 @@ def span_embeddings_forward(
out = model.ops.alloc2f(len(indoc), dim)
for ii, (start, end) in enumerate(dY.indices[offset:hi]):
# adjust indexes to align with doc
start -= tokoffset
end -= tokoffset
out[start] += starts[ii]
out[end] += ends[ii]
out[start:end] += spanvecs[ii]
idxs = dY.indices[offset:hi] - tokoffset
ops.scatter_add(out, idxs[:, 0], starts)
ops.scatter_add(out, idxs[:, 1], ends)
ops.scatter_add(out, idxs.T, spanvecs)
oweights.append(out)
offset = hi