Implement coalesced pooling over entire batches

spacy_transformers/layers/trfs2arrays.py

@@ -1,4 +1,5 @@
-from typing import List, cast
+from typing import Callable, List, Optional, Tuple, cast
+import numpy
 from spacy.util import all_equal
 from transformers.file_utils import ModelOutput
 from transformers.modeling_outputs import BaseModelOutput
@@ -13,15 +14,19 @@ def trfs2arrays(
 ) -> Model[List[TransformerData], List[Floats2d]]:
  """Pool transformer data into token-aligned tensors."""
  return Model(
- "trfs2arrays", forward, layers=[pooling], attrs={"grad_factor": grad_factor}
+ "trfs2arrays",
+ forward,
+ layers=[pooling],
+ attrs={"grad_factor": grad_factor},
  )
 
 
 def forward(model: Model, trf_datas: List[TransformerData], is_train: bool):
  pooling: Model[Ragged, Floats2d] = model.layers[0]
  grad_factor = model.attrs["grad_factor"]
- outputs = []
- backprops = []
+ zero_outputs: List[Tuple[int, Floats2d]] = []
+ backprops_alignment: List[Optional[Callable]] = []
+ aligned_outputs: List[Tuple[int, Ragged]] = []
 
  # For zero-length documents, we could cache the output width by iterating
  # through the batch outputs and retrieving the shape of a non-zero length
@@ -31,55 +36,70 @@ def forward(model: Model, trf_datas: List[TransformerData], is_train: bool):
  # zero in these cases as the effective length of the resultant tensor is zero anyway.
  output_width = 0
 
- for trf_data in trf_datas:
+ for i, trf_data in enumerate(trf_datas):
  if "last_hidden_state" in trf_data.model_output:
  tensor_t_i = cast(BaseModelOutput, trf_data.model_output).last_hidden_state
  if tensor_t_i.size == 0:
  # This can happen during prediction/initialization if the transformer pipe was disabled/not executed and one of the inputs
  # was of length zero. This causes the listenener to generate a zero-sized (in the sequence length dim) TransformerData
  # output and pass it downstream.
- outputs.append(model.ops.alloc2f(0, output_width))
- backprops.append((None, None))
+ zero_outputs.append((i, model.ops.alloc2f(0, output_width)))
+ backprops_alignment.append(None)
  else:
  # This is the general case for non-zero length documents.
  src = model.ops.reshape2f(tensor_t_i, -1, trf_data.width) # type: ignore
  dst, get_d_src = apply_alignment(model.ops, trf_data.align, src)
- output, get_d_dst = pooling(dst, is_train)
- outputs.append(output)
- backprops.append((get_d_dst, get_d_src)) # type: ignore
+ aligned_outputs.append((i, dst))
+ backprops_alignment.append(get_d_src)
  else:
  # This can happen during prediction/training for zero-length documents. Since zero-length docs
  # are implicitly ignored in the span generation stage, the transformer model does not return any
  # predictions for them and subsequently, FullTransformerBatch.split_by_doc() generates an empty
  # TransformerData.
- outputs.append(model.ops.alloc2f(0, output_width))
- backprops.append((None, None))
+ zero_outputs.append((i, model.ops.alloc2f(0, output_width)))
+ backprops_alignment.append(None)
+
+ pooling_outputs, backprop_pooling = concat_pooling_forward(
+ pooling, [dst for _, dst in aligned_outputs], is_train
+ )
+
+ # Interleave the zero and non-zero outputs into the final result.
+ outputs: List[Optional[Floats2d]] = [None] * (
+ len(zero_outputs) + len(aligned_outputs)
+ )
+ for i, zero_output in zero_outputs:
+ outputs[i] = zero_output
+ for (i, _), pooling_output in zip(aligned_outputs, pooling_outputs):
+ outputs[i] = pooling_output
 
  def backprop_trf_to_tensor(d_outputs: List[Floats2d]) -> List[TransformerData]:
  d_trf_datas: List[TransformerData] = []
- to_zip = (trf_datas, d_outputs, backprops)
+
+ # Only update the gradients that are relevant for pooling.
+ d_pooling = backprop_pooling([d_outputs[i] for i, _ in aligned_outputs])
+ for (i, _), d_pooling_i in zip(aligned_outputs, d_pooling):
+ d_outputs[i] = d_pooling_i
+
+ to_zip = (trf_datas, d_outputs, backprops_alignment)
  assert all_equal(len(x) for x in to_zip) # type: ignore
  zipped = zip(*to_zip)
- for trf_data, d_output, (get_d_dst, get_d_src) in zipped:
+ for trf_data, d_output, get_d_src in zipped:
  if "last_hidden_state" not in trf_data.model_output:
  # This gradient belongs to a zero-length doc and must be ignored as it doesn't have a corresponding
  # output from the transformer model (due to empty documents being skipped during the span generation
  # stage in the forward pass).
  assert len(d_output) == 0
  assert get_d_src is None
- assert get_d_dst is None
  continue
 
  assert get_d_src is not None
- assert get_d_dst is not None
  d_model_output = ModelOutput(
  last_hidden_state=model.ops.alloc(
  trf_data.model_output.last_hidden_state.shape, # type: ignore
  dtype=trf_data.model_output.last_hidden_state.dtype, # type: ignore
  )
  )
- d_dst = cast(Floats2d, get_d_dst(d_output))
- d_src = cast(Floats2d, get_d_src(d_dst))
+ d_src = get_d_src(d_output)
  d_src *= grad_factor
  d_model_output["last_hidden_state"] = d_src.reshape(
  cast(BaseModelOutput, trf_data.model_output).last_hidden_state.shape
@@ -95,3 +115,35 @@ def backprop_trf_to_tensor(d_outputs: List[Floats2d]) -> List[TransformerData]:
 
  assert len(outputs) == len(trf_datas)
  return outputs, backprop_trf_to_tensor
+
+
+def concat_pooling_forward(
+ pooling: Model[Ragged, Floats2d], X: List[Ragged], is_train: bool
+):
+ xp = pooling.ops.xp
+
+ datas = []
+ lens = []
+ doc_lens = []
+ for X_doc_data in X:
+ datas.append(X_doc_data.dataXd)
+ lens.append(X_doc_data.lengths)
+ doc_lens.append(len(X_doc_data.lengths))
+
+ X_flat = Ragged(xp.concatenate(datas, axis=0), xp.concatenate(lens, axis=0))
+ Y_pooled, pooling_backprop = pooling(X_flat, is_train)
+ Y = xp.split(Y_pooled, numpy.cumsum(doc_lens)[:-1])
+
+ def backprop(dY):
+ dY_pooled_flat = xp.concatenate(dY)
+ dY_flat = pooling_backprop(dY_pooled_flat).dataXd
+
+ dY = []
+ for X_doc_data in X:
+ doc_unpooled_len = X_doc_data.dataXd.shape[0]
+ dY.append(Ragged(dY_flat[:doc_unpooled_len], X_doc_data.lengths))
+ dY_flat = dY_flat[doc_unpooled_len:]
+
+ return dY
+
+ return Y, backprop