Loading Python model into nn::Module in libtorch

botelho: <p>Is there a way to export a model to file in python and load it into a <code>torch::nn::Module</code> in libtorch?</p>
<p>I know it can be loaded into a <code>jit::script::Module</code>, but that does not solve my problem. The <code>torch::jit</code> API is completely orthogonal to <code>torch::nn</code>, and I do need to use <code>torch::nn::Module</code> only. Thanks!</p>

dfalbel: <p>I don’t think there’s a way to do it.</p>
<p>You can recreate the architecture using <code>torch::nn::Module</code> API and then reload the <code>state_dict()</code> saved from Python though.</p>

botelho: <p>Valeu, <a class="mention" href="/u/dfalbel">@dfalbel</a>!<br>
Thanks for your response. And how do you suggest I extract the architecture from the model saved in python? Assuming all I get is the <code>.pt</code> file that someone saved in python, how can I recreate the architecture in libtorch with only that?</p>

dfalbel: <p>Boa <a class="mention" href="/u/botelho">@botelho</a> !</p>
<p>I don’t think that’s possible. You will need to reimplement the architecture manually and then you can load the weights.<br>
The pt file saved with <code>torch.save</code> contains pickled python objects and needs a Python runtime to be executed.</p>

botelho: <p>Obrigado de novo, <a class="mention" href="/u/dfalbel">@dfalbel</a>!</p>
<p>This really blows my mind! I work on a libtorch-based application in which I can load HDF5 models from Keras. Since the HDF5 contains both the network graph (as a json string) and the weights, I can programmatically construct the <code>nn::Module</code> and populate its trainable parameters. All without going through any format conversion! And, ironically, I cannot do the same with a model saved from <strong>Pytorch itself</strong> in python!! I really wish some Pytorch developer explained the rationale behind making the <code>.pt</code> file so convoluted, and then developing a whole new API (<code>jit::script</code>) just to handle it in libtorch.</p>

dfalbel: <p>I can’t speak for PyTorch devs but I think the json string representation is convenient for this kind of application but also has its drawbacks, for example, if you have custom layers, and functions in your model, you need to recreate them (even in Python) before being able to reload the model. That doesn’t happen with <code>torch.save()</code> because it pickles everything that is needed to reload the model.</p>
<p>What exactly do you need from the <code>torch::nn::Module</code> interface that you don’t have from <code>torch::jit::Module</code>s?</p>

botelho: <p>Thanks again, <a class="mention" href="/u/dfalbel">@dfalbel</a>. You make some good points about the limitations of the Keras approach, which might explain why Pytorch decided to go with the pickle-based solution. Although I suspect Keras has come up with ways to support loading custom layers (via lambdas). Besides, I would be more than happy if Pytorch allowed me to load at least the standard (non-custom) layers directly from python to <code>nn::Module</code>!</p>
<p>As per your question, the main reason I need <code>nn::Module</code> is because we already have an entire application developed around it, and switching to another API is not really feasible. And even if it were, there are some critical use cases that I believe would not be possible. For example, I might want to augment a model by adding more layers and only loading the weights for the old layers (e.g., for transfer-learning). I would not be able to do that with <code>jit::script::Module</code>, since the layers themselves are wrappers around <code>nn::Module</code>.</p>

dfalbel: <aside class="quote no-group" data-username="botelho" data-post="7" data-topic="154159">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/botelho/48/12582_2.png" class="avatar"> botelho:</div>
<blockquote>
<p>Although I suspect Keras has come up with ways to support loading custom layers (via lambdas)</p>
</blockquote>
</aside>
<p>AFAICT the way it works now is by serializing the TensorFlow graph, which is similar to what TorchScript does.</p>
<aside class="quote no-group" data-username="botelho" data-post="7" data-topic="154159">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/botelho/48/12582_2.png" class="avatar"> botelho:</div>
<blockquote>
<p><code>nn::Module</code> is because we already have an entire application developed around it, and switching to another API is not really feasible</p>
</blockquote>
</aside>
<p>Can you perhaps create a <code>torch::nn::Module</code> that has your loaded <code>torch::jit::Module</code> as a submodule? Then you can treat it as any other <code>torch::nn::Module</code>, append layers and etc?</p>

botelho: <aside class="quote no-group" data-username="dfalbel" data-post="8" data-topic="154159">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/dfalbel/48/8880_2.png" class="avatar"> dfalbel:</div>
<blockquote>
<p>Can you perhaps create a <code>torch::nn::Module</code> that has your loaded <code>torch::jit::Module</code> as a submodule</p>
</blockquote>
</aside>
<p>I don’t think so, because only other <code>nn::Modules</code> or their derived types or can be registered as submodules.</p>

dfalbel: <aside class="quote no-group" data-username="botelho" data-post="9" data-topic="154159">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/botelho/48/12582_2.png" class="avatar"> botelho:</div>
<blockquote>
<p>I don’t think so, because only other <code>nn::Modules</code> or their derived types or can be registered as submodules.</p>
</blockquote>
</aside>
<p>FWIW this seems to work for me:</p>
<pre><code class="lang-auto">#include &lt;torch/torch.h&gt;
#include &lt;torch/script.h&gt;
#include &lt;iostream&gt;

int main() {

  // Define a new Module.
  struct Net : torch::nn::Module {
    torch::jit::Module module_;
    
    Net(std::string path) {
      module_ = torch::jit::load(path);
      for (const auto &amp;par : module_.named_parameters()) {
        register_parameter(par.name, par.value);
      }
    }

    // Implement the Net's algorithm.
    torch::Tensor forward(torch::Tensor x) {
      std::vector&lt;torch::jit::IValue&gt; inputs;
      inputs.push_back(x);
      return module_.forward(inputs).toTensor();
    }

  };

  auto model = Net("./linear.pt");
  auto x = torch::ones({10, 10});
  auto y = (2*model.forward(x)).sum();
  y.backward();
  std::cout &lt;&lt; model.parameters()[0].grad() &lt;&lt; std::endl;
  
  return 0;
}
</code></pre>
<p>So in theory you can wrap the jit module into a <code>nn::Module</code> and then call forward, backprop into it, and etc.</p>

botelho: <p>Thank you, that’s a promising idea! I had thought in terms of actual torch submodules, as in calling <code>register_module</code> on the <code>nn::Module</code>. But wrapping the <code>jit::Module</code> as member should work too, I’ll see where that takes me. There’s some immediate kinks to iron out, e.g.,</p>
<pre><code class="lang-auto">terminate called after throwing an instance of 'c10::Error'
  what():  Parameter name must not contain a dot (got 'conv1.weight')
</code></pre>
<p>It seems like registering the parameters will be a little more complicated…</p>

dfalbel: <aside class="quote no-group" data-username="botelho" data-post="11" data-topic="154159">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/botelho/48/12582_2.png" class="avatar"> botelho:</div>
<blockquote>
<pre><code class="lang-auto">terminate called after throwing an instance of 'c10::Error'
  what():  Parameter name must not contain a dot (got 'conv1.weight')
</code></pre>
</blockquote>
</aside>
<p>Yeah, I guess you can recurse, eg with something like:</p>
<pre><code class="lang-auto">Net(torch::jit::Module module) {
      module_ = module;

      for (const auto &amp;mod : module_.named_children()) {
        register_module(mod.name, std::make_shared&lt;Net&gt;(mod.value));
      }

      for (const auto &amp;par : module_.named_parameters(false)) {
        register_parameter(par.name, par.value);
      }
    }
</code></pre>

botelho: <p>Yep, that should work. Thanks again!</p>