Load methods via get_method from torchscript module

[awvwgk]

For an application we store a torchscript checkpoint with multiple entry points, the forward entry point is considered not part of our public API since it requires some extra processing. Therefore, we export our checkpoint with extra methods and load those from the module via get_method.

Here is an example on how we do this in libtorch

#include <torch/script.h>

torch::jit::Method
load_method(const std::string &filename, torch::DeviceType device)
{
  torch::jit::script::Module mod;
  torch::jit::ExtraFilesMap extra_files;

  try
  {
    // Deserialize the ScriptModule from a file using torch::jit::load().
    mod = torch::jit::load(filename, device, extra_files);
  }
  catch (const c10::Error &e)
  {
    throw std::runtime_error("Error loading the model from " + filename + ": " + e.what());
  }

  return mod.get_method("get_exc_density");
}

Is this possible with the torch_model type in FTorch as well?

[jatkinson1000]

Hi @awvwgk this is not currently a feature we have implemented, but looking at the sample code above it seems like it should be relatively straightforward. It seems that we would need to:

• Provide an additional optional argument of extra_files which would be a list of character strings on the Fortran side?
• Pass this through to the jit::load call.

Does this sound correct?
Would anything further be expected/required?

If so we should be able to look at this, please open an issue with it as a feature request referencing this discussion.

What would be useful is if you could provide us with a MWE from the PyTorch side that we could use to build and test against as we're not super familiar with this feature.

CC: @joewallwork

[awvwgk]

I can give a try to implement this feature in FTorch.

[jatkinson1000]

@awvwgk we'd welcome the contribution and are happy to support external contributors with understanding the codebase etc.

Does my summary above correctly capture what you're after or is there anything else to be added?

Once you have something (even if it is work in progress) feel free to open a Pull Request and we can review/provide feedback.

[awvwgk]

I opened a draft in #539.

[jatkinson1000]

A quick thought (having now read #536) would be that the Fortran extra_files object would need to be unpacked into a torch::jit::ExtraFilesMap. It may be that it is possible to do this within C++ from an array of strings, but if the IDs of the map are also required it may be that wrapping the torch::jit::ExtraFilesMap type is also of use (though that might come after proof-of concept and a simple example 😉).

[jatkinson1000]

Thanks, We'll try and take a look at that this week!