dotnet / dotnet/machinelearning

Reduce steady-state allocations in inference and training hot paths

Open
#7,684 2 comments 0 reactions 0 assignees View on GitHub
area-Core tenet-performance
Dominant language
C#
Stars
9.4k
Forks
2k
Avg merge
2d 20h
Merged PRs (30d)
11

Description

A source-level audit found several avoidable managed allocations in ML.NET inference, training, and shared data-pipeline hot paths. These occur per row, example, class, or training iteration rather than only during model or pipeline initialisation.

| Area | Evidence | Concern |
|---|---|---|
| TensorFlow inference | `src/Microsoft.ML.TensorFlow/TensorflowTransform.cs:756,903`, `TensorflowUtils.cs:447` | Creates a `Runner`, input/output arrays, shape array, and full input buffer per row. |
| ONNX inference | `src/Microsoft.ML.OnnxTransformer/OnnxTransform.cs:689`, `OnnxUtils.cs:559,570` | Copies each input and shape into new arrays and creates a new input list per row. |
| SDCA training | `src/Microsoft.ML.StandardTrainers/Standard/SdcaBinary.cs:847`, `SdcaMulticlass.cs:265` | Capturing lambdas allocate inside the innermost compare-and-swap loop. Multiclass multiplies this by class count and retry count. |
| Shared vector processing | `src/Microsoft.ML.DataView/VBuffer.cs:342,471` | `VBuffer.Items()` uses a `yield` iterator, producing an allocation for each enumeration. Hot callers include FFM, FastTree and prediction materialisation. |
| Multiclass inference | `src/Microsoft.ML.StandardTrainers/Standard/MulticlassClassification/OneVersusAllTrainer.cs:636,729,885,897` | Invokes `Parallel.For` with a captured delegate per row. Softmax also creates a `double[classCount]` array per prediction. |
| Prediction output | `src/Microsoft.ML.Data/DataView/TypedCursor.cs:459-469` | Sparse vector outputs allocate a new array even when using the reusable `PredictionEngine.Predict(..., ref prediction)` overload. |
| Text processing | `src/Microsoft.ML.Transforms/Text/NgramTransform.cs:698,715` | TF-IDF and IDF getters construct capturing delegates per row. |
| Cursor consolidation | `src/Microsoft.ML.Data/Data/DataViewUtils.cs:406,728` | Creates a `Batch` and one `BatchColumn` wrapper per active column at every batch boundary. |
| Enumerable ingestion | `src/Microsoft.ML.Data/DataView/DataViewConstructionUtils.cs:395,400` | Key conversion through `Convert.ChangeType` boxes values twice per row. |

The existing performance suite does not adequately detect these regressions:

- `RecommendedConfig` does not enable BenchmarkDotNet's `MemoryDiagnoser`.
- `BenchmarksTest` checks only that benchmarks build and execute, not allocation thresholds.
- The batch inference benchmarks at `StochasticDualCoordinateAscentClassifierBench.cs:166-172` call the lazy `Transform` method without enumerating its output, so they do not measure scoring.

**Describe the solution you'd like**

Add allocation-focused benchmarks and remove the highest-frequency allocations while preserving numerical behaviour and existing concurrency guarantees.

- [ ] Enable allocation reporting for relevant performance benchmarks.
- [ ] Add separate inference benchmarks for allocating and reusable `PredictionEngine` overloads, including dense and sparse vector outputs.
- [ ] Add fixed-shape ONNX and TensorFlow per-row inference benchmarks.
- [ ] Add controlled binary and multiclass SDCA training benchmarks.
- [ ] Ensure batch inference benchmarks enumerate the resulting `IDataView`.
- [ ] Cache TensorFlow runners, input/output names and shape arrays where their lifetimes permit.
- [ ] Reuse TensorFlow and ONNX fixed-size input buffers rather than copying them into new managed arrays per row.
- [ ] Replace the SDCA visitor delegate with an atomic operation on the dual-table abstraction.
- [ ] Replace hot internal `VBuffer.Items()` calls with span/index-based loops or an allocation-free enumerator.
- [ ] Reuse the OVA softmax scratch buffer and avoid `Parallel.For` below a measured class-count threshold.
- [ ] Reuse same-sized sparse prediction arrays through `VBuffer.CopyTo(Span)`.
- [ ] Construct n-gram weighting delegates once per getter rather than once per row.
- [ ] Investigate pooling the lightweight cursor-consolidation wrappers.
- [ ] Replace per-row `Convert.ChangeType` calls with converters selected during getter construction.

Acceptance criteria:

- Allocation and throughput results are recorded before and after each change.
- Fixed-shape ONNX and TensorFlow inference no longer copies full input vectors into new managed arrays on every row where backend lifetime rules permit reuse.
- SDCA performs no delegate or closure allocation per dual update.
- Targeted FFM, FastTree and prediction materialisation paths do not allocate iterator objects per row.
- Reused sparse prediction outputs do not allocate a replacement array when their length is unchanged.
- Numerical output and deterministic single-threaded training behaviour remain unchanged.
- Buffers remain scoped per mapper, cursor, or prediction engine where required for thread safety.

**Describe alternatives you've considered**

Relying on garbage collection or GC tuning would only mitigate the symptoms and would not remove allocation volume from the hottest loops.

Advising callers to use `PredictionEngine.Predict(..., ref prediction)` is useful, but it addresses only the result-object allocation. It does not fix tensor copies, sparse output arrays, multiclass scoring temporaries, or training allocations.

Global shared buffer pools were also considered, but per-engine or per-cursor reuse is safer because `PredictionEngine` and several mapper caches are intentionally not thread-safe.

**Additional context**

ML.NET already has suitable allocation-conscious patterns, particularly `VBufferEditor.Create(ref destination, ...)`, reusable getter buffers, and preallocated trainer scratch arrays. The proposed changes should extend those existing patterns rather than introduce broadly shared mutable state.

This issue is based on static source inspection. Runtime profiling and allocation baselines should be added before implementation to quantify each finding and prioritise the work.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.