loadModel of a simple LSTM model, with weights on the order of 3MB uses 5GB memory.

Code to reproduce:

const model = await tf.loadModel('https://timotheebernard.github.io/models/model.json');

// This yields:
// { 
//    unreliable: false
//    numTensors: 8991
//    numDataBuffers: 8986
//    numBytes: 5050462344
// }
console.log(tf.memory());

// If you use this code here you can get the total size of the weights of all layers,
//  which gives 766601. With 4 bytes per weight, that's 3066404 bytes ~= 3MB.
const totalWeightsSize =
    tf.util.flatten(model.layers.map(l => l.weights).filter(x => x.length > 0))
    .map(x => x.val)
    .reduce((accumulator, weight) => accumulator += weight.size, 0);

Note that the weights files are small, < 3MB: https://github.com/timotheebernard/timotheebernard.github.io/tree/master/models

Is there something that recurrent cells are doing that cause this much memory blowout? @caisq, @ericdnielsen, @bileschi can you take a look?

[caisq]

My guess is that this has to do with the Orthogonal initializer of the LSTM. When the model JSON is loaded, the Orthogonal initializer calls the QR decomposition under the hood. Owing to the relatively large size of the matrix, this slows things done. Maybe we can optimize the memory usage of the QR decomposition a little. We can also consider adding the logic to skip any initialization if weights are going to be loaded afterwards.

[caisq]

FYI, I plan to speed up orthogonal initializers by replacing the QR decomposition with the Gram-Schmidt process.