Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support loading large model weights #7610

Merged
merged 9 commits into from
Apr 24, 2023
Merged

Support loading large model weights #7610

merged 9 commits into from
Apr 24, 2023

Commits on Apr 20, 2023

  1. webgpu: Fix a bug in softmax (tensorflow#7607)

    @qjia7
    qjia7 authored Apr 20, 2023
    Configuration menu
    Copy the full SHA
    78a3a03 View commit details
    Browse the repository at this point in the history

  2. Avoid allocating a large arraybuffer when loading weights (tensorflow… 

    …#7598)

The loadWeights function loads weights in 4MB chunks and then concatenates them into a single large ArrayBuffer. That ArrayBuffer is used for splitting the weights data back up into tensors. Allocating large ArrayBuffers (3.5GB) can be unstable on Chrome, so this PR avoids this allocation, instead slicing the weights out of the chunks manually.

The implementation wraps the array of weights (stored as ArrayBuffer[]) in a new CompositeArrayBuffer class. This class implements slice by copying the desired range out of the buffer(s) that it overlaps with.
    @mattsoulanille
    mattsoulanille authored Apr 20, 2023
    Configuration menu
    Copy the full SHA
    3ceace9 View commit details
    Browse the repository at this point in the history

  3. Support using a list of ArrayBuffers as model weight data

    @mattsoulanille
    mattsoulanille committed Apr 20, 2023
    Configuration menu
    Copy the full SHA
    cf2493e View commit details
    Browse the repository at this point in the history

  4. Avoid 'Array.flat()'

    @mattsoulanille
    mattsoulanille committed Apr 20, 2023
    Configuration menu
    Copy the full SHA
    6720fde View commit details
    Browse the repository at this point in the history

  5. Simplify some of the tests

    @mattsoulanille
    mattsoulanille committed Apr 20, 2023
    Configuration menu
    Copy the full SHA
    4d22064 View commit details
    Browse the repository at this point in the history

  6. Do not export 'CompositeArrayBuffer' from tfjs-core

    @mattsoulanille
    mattsoulanille committed Apr 20, 2023
    Configuration menu
    Copy the full SHA
    b21b302 View commit details
    Browse the repository at this point in the history

  7. Update doc for weightData

    @mattsoulanille
    mattsoulanille committed Apr 20, 2023
    Configuration menu
    Copy the full SHA
    775e31d View commit details
    Browse the repository at this point in the history

  8. Fix tfjs-node

    @mattsoulanille
    mattsoulanille committed Apr 20, 2023
    Configuration menu
    Copy the full SHA
    d2613b7 View commit details
    Browse the repository at this point in the history

  9. Remove unused import

    @mattsoulanille
    mattsoulanille committed Apr 20, 2023
    Configuration menu
    Copy the full SHA
    0276e01 View commit details
    Browse the repository at this point in the history