TorchScript Export

[mvidela31]

Hi everybody,

I was wondering if it is possible to export Anomalib's Pytorch Lighting models to TorchScript format. I tried to use the to_torchscript() method as suggested in the Pytorch lightning docs on the PatchCore model with no success:

import torch
from anomalib.models import get_model
from anomalib.config import get_configurable_parameters

# Artifact paths
model_path = './results/patchcore/mvtec/bottle/weights/model.ckpt'
config_path = './anomalib/models/patchcore/config.yaml'
# Read config file
cfg = get_configurable_parameters(config_path=config_path)
# Load model
model = get_model(cfg)
model.load_state_dict(torch.load(model_path)["state_dict"])
model.eval()
# Compile the Pytorch lightning model to TorchScript
script = model.to_torchscript()
# Save for use in production environment
torch.jit.save(script, 'patchcore.pt')

---------------------------------------------------------------------------

RuntimeError                              Traceback (most recent call last)

[<ipython-input-7-069fb611dfec>](https://localhost:8080/#) in <module>
----> 1 script = model.to_torchscript()
      2 # save for use in production environment
      3 torch.jit.save(script, "patchcore.pt")

11 frames

[/usr/local/lib/python3.7/dist-packages/torch/autograd/grad_mode.py](https://localhost:8080/#) in decorate_context(*args, **kwargs)
     25         def decorate_context(*args, **kwargs):
     26             with self.clone():
---> 27                 return func(*args, **kwargs)
     28         return cast(F, decorate_context)
     29 

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/core/lightning.py](https://localhost:8080/#) in to_torchscript(self, file_path, method, example_inputs, **kwargs)
   1945 
   1946         if method == "script":
-> 1947             torchscript_module = torch.jit.script(self.eval(), **kwargs)
   1948         elif method == "trace":
   1949             # if no example inputs are provided, try to see if model has example_input_array set

[/usr/local/lib/python3.7/dist-packages/torch/jit/_script.py](https://localhost:8080/#) in script(obj, optimize, _frames_up, _rcb, example_inputs)
   1264         obj = call_prepare_scriptable_func(obj)
   1265         return torch.jit._recursive.create_script_module(
-> 1266             obj, torch.jit._recursive.infer_methods_to_compile
   1267         )
   1268 

[/usr/local/lib/python3.7/dist-packages/torch/jit/_recursive.py](https://localhost:8080/#) in create_script_module(nn_module, stubs_fn, share_types, is_tracing)
    452     if not is_tracing:
    453         AttributeTypeIsSupportedChecker().check(nn_module)
--> 454     return create_script_module_impl(nn_module, concrete_type, stubs_fn)
    455 
    456 def create_script_module_impl(nn_module, concrete_type, stubs_fn):

[/usr/local/lib/python3.7/dist-packages/torch/jit/_recursive.py](https://localhost:8080/#) in create_script_module_impl(nn_module, concrete_type, stubs_fn)
    514 
    515     # Actually create the ScriptModule, initializing it with the function we just defined
--> 516     script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
    517 
    518     # Compile methods if necessary

[/usr/local/lib/python3.7/dist-packages/torch/jit/_script.py](https://localhost:8080/#) in _construct(cpp_module, init_fn)
    592             """
    593             script_module = RecursiveScriptModule(cpp_module)
--> 594             init_fn(script_module)
    595 
    596             # Finalize the ScriptModule: replace the nn.Module state with our

[/usr/local/lib/python3.7/dist-packages/torch/jit/_recursive.py](https://localhost:8080/#) in init_fn(script_module)
    492             else:
    493                 # always reuse the provided stubs_fn to infer the methods to compile
--> 494                 scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
    495 
    496             cpp_module.setattr(name, scripted)

[/usr/local/lib/python3.7/dist-packages/torch/jit/_recursive.py](https://localhost:8080/#) in create_script_module_impl(nn_module, concrete_type, stubs_fn)
    514 
    515     # Actually create the ScriptModule, initializing it with the function we just defined
--> 516     script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
    517 
    518     # Compile methods if necessary

[/usr/local/lib/python3.7/dist-packages/torch/jit/_script.py](https://localhost:8080/#) in _construct(cpp_module, init_fn)
    592             """
    593             script_module = RecursiveScriptModule(cpp_module)
--> 594             init_fn(script_module)
    595 
    596             # Finalize the ScriptModule: replace the nn.Module state with our

[/usr/local/lib/python3.7/dist-packages/torch/jit/_recursive.py](https://localhost:8080/#) in init_fn(script_module)
    492             else:
    493                 # always reuse the provided stubs_fn to infer the methods to compile
--> 494                 scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
    495 
    496             cpp_module.setattr(name, scripted)

[/usr/local/lib/python3.7/dist-packages/torch/jit/_recursive.py](https://localhost:8080/#) in create_script_module_impl(nn_module, concrete_type, stubs_fn)
    518     # Compile methods if necessary
    519     if concrete_type not in concrete_type_store.methods_compiled:
--> 520         create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
    521         # Create hooks after methods to ensure no name collisions between hooks and methods.
    522         # If done before, hooks can overshadow methods that aren't exported.

[/usr/local/lib/python3.7/dist-packages/torch/jit/_recursive.py](https://localhost:8080/#) in create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
    369     property_rcbs = [p.resolution_callback for p in property_stubs]
    370 
--> 371     concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
    372 
    373 def create_hooks_from_stubs(concrete_type, hook_stubs, pre_hook_stubs):

RuntimeError: 
Module 'FeatureExtractor' has no attribute 'layers' (This attribute exists on the Python module, but we failed to convert Python type: 'omegaconf.listconfig.ListConfig' to a TorchScript type. Only tensors and (possibly nested) tuples of tensors, lists, or dictsare supported as inputs or outputs of traced functions, but instead got value of type ListConfig.. Its type was inferred; try adding a type annotation for the attribute.):
  File "/content/anomalib/anomalib/models/components/feature_extractors/feature_extractor.py", line 88
            Feature map extracted from the CNN
        """
        features = dict(zip(self.layers, self.feature_extractor(input_tensor)))
                            ~~~~~~~~~~~ <--- HERE
        return features

It is possible to export the Anomalib's models to TorchScript format? If not, please consider adding this useful feature in the future.

Many thanks.

[andreluizbvs]

From my experience with TorchScript and Anomaly Detection models, it is really not that simple (e.g., just call torch.jit.save()). In my case, depending on the cues given by the compiler, I would try to untangle the problematic line of code to get to the core of the issue. Sometimes, just by doing this, the problem would go away.

So, in your case, I would try to untagle the line features = dict(zip(self.layers, self.feature_extractor(input_tensor))) into multiple lines. Notice that there is a bunch of stuff happening in only one line (e.g., zip and dictionary conversion/instantiation, feature_extractor call, etc.). Try doing something like that below to get to the core of the problem:

features = self.feature_extractor(input_tensor)
features = zip(self.layers, features)
features = dict(features)

That helped me quite a lot during my debugging phase of jit compilation of torch models.

[Yizhang-Li]

save issue here

[ashwinvaidya17]

@mvidela31 Anomalib currently supports OpenVINO export (including patchcore). This enables faster inference speed and allows edge deployment in low power devices. With OpenVINO you can also use these models from your C++ code.
While we currently do not have a plan to support TorchScript, can you tell us what specific features are you looking for from TorchScript models. We can then consider introducing it to Anomalib.

[ashwinvaidya17]

@Yizhang-Li Correct me if I am wrong but you will have to write an inference script which uses the jit exported model. Can't you run a script similar to our openvino_inference.py?

[Yizhang-Li]

@Yizhang-Li Correct me if I am wrong but you will have to write an inference script which uses the jit exported model. Can't you run a script similar to our openvino_inference.py?

Yes. I was planning to write the inference script. Thanks for the suggestion. I will try it!

[mvidela31]

Hi @ashwinvaidya17,

I was planning to deploy the Anomalib models on AWS using the Amazon EC2 Inf1 instance (recommended for efficient and low-cost inference of Deep Learning models). As you can see, this inference efficiency relies on compiling the models on Inferentia hardware using the AWS Neuron SDK. However, the AWS Neuron compilation method only supports Tensorflow, Pytorch, and MXNet models. Therefore, being able to export the Anomalib models to TorchScript format would allow them to be compiled on Inferentia hardware using the AWS Neuron SDK (as confirmed here), extending them usability on the AWS platform.

[ashwinvaidya17]

@mvidela31 Thanks for letting us know. We will keep this in mind for future support.

[vmiller987]

@mvidela31 Thanks for letting us know. We will keep this in mind for future support.

I just wanted to add I wouldn't mind TorchScript support as well, however, I plan to look into OpenVINO and maybe convince my team to swap off TorchScript. I've ran into other issues with TorchScript, so I'm not a fan. The current next lib they would prefer is ONNX.