Importing models requiring other Torch modules

[msjgriffiths]

I have a model trained internally with two extensions:

• torchtext: Using SentencePiece.
• sru: A C++ implementation of a custom RNN unit.

I can try to load the model in R with:

library(torchtext)

m1 <- jit_load("RecursiveScriptModule_model.zip")

However, that provides an error:

Error in cpp_jit_load(path) : 
Unknown type name '__torch__.torch.classes.torchtext.SentencePiece':
Serialized   File "code/__torch__/asapp/valley/autosuggest/autosuggest_new/inference.py", line 5
  __buffers__ = []
  training : bool
  tokenizer : __torch__.torch.classes.torchtext.SentencePiece
              ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
  whitespace_id : int
  context_len : int

I can replicate this error in Python by not loading torchtext, e.g. this provides the same error:

import torch
model = torch.jit.load("RecursiveScriptModule_model.zip") # Fails with same error

Instead I need to load the dependencies in Python, e.g.

import torch, torchtext, sru
model = torch.jit.load("RecursiveScriptModule_model.zip") # Works!

I can also load this in C++. In the CMakeLists.txt file, I have:

set(Sentencepiece_ROOT ${CMAKE_SOURCE_DIR}/third_party)
find_package(Sentencepiece)
if(NOT Sentencepiece_FOUND)
    message(STATUS "Installing Sentencepiece...")
    execute_process(COMMAND
            ${CMAKE_CURRENT_LIST_DIR}/tools/sentencepiece_install.sh
            ${CMAKE_SOURCE_DIR}/third_party)
    find_package(Sentencepiece REQUIRED)
endif()

add_subdirectory(sru/sru/csrc)

I would like to add the same dependencies when loading in R. Since I know the C++ dependencies, can I pass them to R? Do I need to make some R package that provides the C++ dependencies? How do I register that?

[dfalbel]

Hey @msjgriffiths

Making it easier to load external c++ extensions is something we want to work soon. I have done a small exploration with torchvision here: mlverse/torchvision#54 (comment)

Basically you need to dyn.load() the extension binary and everything should just work, no need for custom registrations and things like this.

[dfalbel]

This should be possible now with torchexport and related features. See:

• Overview: https://blogs.rstudio.com/ai/posts/2022-04-27-torch-outside-the-box/
• An example: https://github.com/mlverse/lltm
• Torchexport docs: https://github.com/mlverse/torchexport
• A working example: https://github.com/mlverse/torchvisionlib