Is it possible to convert torchnmf model to TorchScript

[dhrubo-os]

I have a use case where I need to convert my learning algorithm (which uses NMF under the hood) into TorchScript. Is it possible to convert torchnmf to torchScript?

I tried with the following example:

import torch
import torch.nn as nn
from torchnmf.nmf import NMF

class Chain(nn.Module):
    def __init__(self):
        super().__init__()
        self.nmf1 = NMF((1024, 10), rank=4)
        self.nmf2 = NMF(W=(24, 10))
        self.nmf3 = NMF(W=(64, 24))

    def forward(self):
        WH = self.nmf1()
        WWH = self.nmf2(H=WH)
        WWWH = self.nmf3(H=WWH)
        return WWWH

model = Chain()
output = model()

jit_processing = torch.jit.script(Chain())
output = jit_processing.forward()
print(output)

print(jit_processing.code)
torch.jit.save(jit_processing, "pytorch_nmf.pt")

But i'm facing error in the line of jit processing : jit_processing = torch.jit.script(Chain())

This is the error log:

Traceback (most recent call last):
  File "/Volumes/workplace/OSML_Metrics_Correlation/src/OSML_Metrics_Correlation/nmf_pytorch.py", line 21, in <module>
    jit_processing = torch.jit.script(Chain())
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_script.py", line 1257, in script
    return torch.jit._recursive.create_script_module(
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_recursive.py", line 451, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_recursive.py", line 513, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_script.py", line 587, in _construct
    init_fn(script_module)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_recursive.py", line 491, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_recursive.py", line 517, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/usr/local/lib/python3.9/site-packages/torch/jit/_recursive.py", line 368, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
RuntimeError: 
Expected a default value of type Tensor on parameter "H".:
  File "/Users/dhrubo/Library/Python/3.9/lib/python/site-packages/torchnmf/nmf.py", line 258
    def forward(self, H: Tensor = None, W: Tensor = None) -> Tensor:
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        r"""An outer wrapper of :meth:`self.reconstruct(H,W) <torchnmf.nmf.BaseComponent.reconstruct>`.
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    
        .. note::
        ~~~~~~~~~
                Should call the :class:`BaseComponent` instance afterwards
                ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
                instead of this since the former takes care of running the
                ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
                registered hooks while the latter silently ignores them.
                ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    
        Args:
        ~~~~~
            H(Tensor, optional): input activation tensor H. If no tensor was given will use :attr:`H` from this module
            ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
                                instead
                                ~~~~~~~
            W(Tensor, optional): input template tensor W. If no tensor was given will use :attr:`W` from this module
            ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
                                instead
                                ~~~~~~~
    
        Returns:
        ~~~~~~~~
            Tensor: tensor
            ~~~~~~~~~~~~~~
        """
        ~~~
        if H is None:
        ~~~~~~~~~~~~~
            H = self.H
            ~~~~~~~~~~
        if W is None:
        ~~~~~~~~~~~~~
            W = self.W
            ~~~~~~~~~~
        assert H is not None
        ~~~~~~~~~~~~~~~~~~~~
        assert W is not None
        ~~~~~~~~~~~~~~~~~~~~
        return self.reconstruct(H, W)
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE


Process finished with exit code 1

[yoyolicoris]

Hi @dhrubo-os ,

Thanks for reporting this. Currently, torchnmf has no test to ensure jittability, but I agree it would be nice to have this feature available.
Will look into it later when I have time.