torch.jit._load_for_mobile with device and extra files options

raedle commented 2 years ago

Summary

The PlayTorch API for torch.jit._loadForMobile did not support arguments for torch device and extra files. This change updates the torch.jit._loadForMobile API to also allow for both parameters. The extra files are useful to store additional data with the model (e.g., store image classes for inference alongside the model).

Changelog

Implemented device arg and extra files arg for torch.jit._loadForMobile and torch.jit._loadForMobileSync (latter is the same but synchronous API).
Added typedoc to torch.ts including a Device type scoped to 'cpu'.
Updated the multi_input_test_modules.py to export extra_files with the model for testing, and re-exported the dummy_test_model.ptl

Test Plan

Example:

let modelUrl = await MobileModel.download(
  require('../../assets/models/dummy_test_model.ptl'),
);

const extraFiles = {foo: undefined, 'model/classes.json': undefined, 'classes.json': undefined};
const model = await torch.jit._loadForMobile<TestModule>(modelUrl, 'cpu', extraFiles);
let lastTime = performance.now();

for (let i = 0; i < 1000; i++) {
  await model.bump(1);
}
let totalTime = performance.now() - lastTime;
console.log('average time for bump: ', totalTime / 1000);
console.log(model);
console.log('extraFiles', extraFiles);

Output:

 LOG  {"forward": [Function forward], "forwardSync": [Function forwardSync]}
 LOG  extraFiles {"classes.json": "", "foo": "", "model/classes.json": "[
  \"Arizona Bark Scorpion\",
  \"Giant Hairy Desert Scorpion\",
  \"Giant Whip Scorpion\",
  \"Northern Scorpion\",
  \"Pseudoscorpiones\",
  \"Stripe-tailed Scorpion\",
  \"Tailless Whip Scorpion\"
]
"}
 LOG  pi:  3.14
 LOG  piSync:  3.14

vercel[bot] commented 2 years ago

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Updated
playtorch	✅ Ready (Inspect)	Visit Preview	Nov 4, 2022 at 11:31PM (UTC)