pytorch model loader

[faroit]

deeplearn.js supported some level of conversion of pytorch models via python scripts

Will this functionality be ported to tfjs or is this out of scope now that the framework is aiming for a tensorflow only solution?

[nsthorat]

We're not planning on supporting pytorch in the immediate future, however there is a path to use the old checkpoint loader with pytorch.

The legacy loader with the pytorch weight-dumper script can be used: https://github.com/PAIR-code/deeplearnjs-legacy-loader

The legacy loader currently depends on deeplearn, so you're going to have a dependency problem for the JS (the python will still work).

Thankfully, all of the logic lives in a single file with 1 dependency (Tensor). This means you could just copy this file into your project: https://github.com/PAIR-code/deeplearnjs-legacy-loader/blob/master/src/checkpoint_loader.ts

You can swap the import of 'deeplearn' with '@tensorflow/tfjs' or '@tensorflow/tfjs-core' and it should just work.

[faroit]

Thanks a lot for the pointer. I will try that.

[slawojstanislawski]

I had a .pth file (containing a pytorch model I presume) and using the method outlined above I was able to get the output: a long list of _weight and _bias files, as well as one manifest.json file. This is in line with what's described here. Finally calling the getAllVariables() method of the legacy loader gets me the variables. However, I can't figure out the steps to move from this vars variable to a usable, tensorflowjs-importable model. Is it by writing out the model manually as it's described on the linked page (I'm not the author of the model, so it's problematic), use tfjs converter in some way, or there's a different procedure?