dotnet / dotnet/machinelearning

Model Saving / Loading memory usage

Open
#905 2 comments 0 reactions 0 assignees View on GitHub
area-Core enhancement Priority:2
Dominant language
C#
Stars
9.4k
Forks
2k
Avg merge
2d 20h
Merged PRs (30d)
11

Description

In reviewing the TensorFlow saving/loading code [PR](https://github.com/dotnet/machinelearning/pull/853) I noticed that we were creating very large byte arrays in the [frozen model case](https://github.com/dotnet/machinelearning/blob/52aff025df29cc02c00999c5ca4a0833a658d142/src/Microsoft.ML.TensorFlow/TensorflowTransform.cs#L123). I believe these can be *very* large models (100MB - many GB) so we might approach the upper limit of the size of managed arrays, not to mention the memory usage of shuttling these bytes into a managed array, just so that they can then be interpreted/read into TensorFlow's object model.

My understanding of the ModelSaveContext / ModelLoadContext get backed with a ZipArchive. This will be backed by a stream (file or memory) and thus not load the entire ML.NET model into memory at one time. As entries are accessed these get loaded as streams which load from disk on demand / decompress on demand. As such the backing stack for model loading/saving permits a minimal memory footprint for loading and saving models. The problem comes in with usage. Many cases where folks are using ReadByteArray they should instead be using a Stream that is constrained to the length of the entry.

I discussed this a bit with @abgoswam and prototyped a sample stream that would wrap the context's BinaryReader/Writer stream and only expose a region. I'm sure we have a better impl floating around somewhere.

I looked a a few of the usage cases of `ReadByteArray` and the following all look suspect of containing large payloads:

https://github.com/dotnet/machinelearning/blob/5e08fa1ea7bfb54f28ed0815cb6413e0068e6dd1/src/Microsoft.ML.Api/SerializableLambdaTransform.cs#L66-L74

I imagine they payload passed to the `LoadDelegate` could be arbitrarily large and should be streamed.

https://github.com/dotnet/machinelearning/blob/307b38f4c86cc31a6a0dbff8c1a302d66f4fe7e7/src/Microsoft.ML.Data/DataLoadSave/PartitionedFileLoader.cs#L226-L234

It looks like the byte array gets stored off in a memory stream until it later gets read as files. I don't see much value in the additional byte-array. Why not instead keep a stream open to the entry in the ModelLoadContext for the lifetime of the loader?

https://github.com/dotnet/machinelearning/blob/5e08fa1ea7bfb54f28ed0815cb6413e0068e6dd1/src/Microsoft.ML.Transforms/OptionalColumnTransform.cs#L99-L103

Similar to above.

https://github.com/dotnet/machinelearning/blob/5e08fa1ea7bfb54f28ed0815cb6413e0068e6dd1/src/Microsoft.ML.Parquet/ParquetLoader.cs#L193-L195

Similar to above.

https://github.com/dotnet/machinelearning/blob/52aff025df29cc02c00999c5ca4a0833a658d142/src/Microsoft.ML.TensorFlow/TensorflowTransform.cs#L123-L141
In the TensorFlowTransform case I don't see a great alternative since the [TF API](https://github.com/tensorflow/tensorflow/blob/3be04971716fcaf0c11ad9262e60efa428553e14/tensorflow/c/c_api.h#L1018-L1020) we call expects a byte buffer and doesn't have a stream-like API for importing the graph (as far as I can tell). It's possible I'm missing something though. A slightly hacky alternative would be to make a memory mapped file, write to that, and pass the [pointer](https://docs.microsoft.com/en-us/dotnet/api/microsoft.win32.safehandles.safememorymappedviewhandle?view=netframework-4.7.2#methods) of the memory mapped file to tensorflow. That way we're never dealing in large byte arrays. We could have some size threshold at which we switch from using a byte array to a memory mapped file. Since there would likely be some tradeoff to creating the MMF.

/cc @Zruty0

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.