Optimize_for_mobile methods from preserved_methods not available in Android/Java via runMethod

username: m-toman
cooked: <p>Hi,</p>
<p>I successfully exported a module with a method:</p>
<pre><code class="lang-auto">    @torch.jit.export
     def inference(self, text_inputs, speaker_embed, style_embed):
</code></pre>
<p>as</p>
<pre><code class="lang-auto">ts_prosody = torch.jit.script(model.prosody_generator)
ts_prosody.save(f"../inference/{model_name}_prosody.pt")
ts_prosody_mobile = optimize_for_mobile(ts_prosody, preserved_methods=["inference", "inference_half"])
torch.jit.save(ts_prosody_mobile,
                      f"../inference/{model_name}_prosody_mobile.pt")
</code></pre>
<p>I verified that in Python both ts_prosody_mobile.inference() and ts_prosody_mobile.inference_half() exist.</p>
<p>Now in Java/Android I load the models with Module.load</p>
<p>then<br>
<code>        IValue[] prosodyOutput = modelProsody.runMethod("inference", ... );</code><br>
works.</p>
<p>But<br>
<code>      IValue[] prosodyOutput = modelProsodyMobile.runMethod("inference", ... );</code></p>
<p>gives me<br>
<code>java.lang.IllegalArgumentException: Undefined method inference</code></p>
<p>Is there anything I missed?</p>
<p>I am using in Python:<br>
torch.<strong>version</strong><br>
‘1.9.0+cu111’</p>
<p>And Java from the Gradle file:<br>
implementation ‘org.pytorch:pytorch_android:1.9.0’</p>
<p>Thanks!</p>

username: guangy10
cooked: <p><a class="mention" href="/u/jacobszwejbka">@JacobSzwejbka</a> could you help with this?</p>

username: m-toman
cooked: <p>Hi,</p>
<p>I’ve tried to work around the issue as:</p>
<pre><code class="lang-auto">class ProsodyGeneratorExport(ProsodyGenerator):
    def __init__(self, hparams):
        super().__init__(hparams)

    def forward(self,
                text_inputs: Tensor,
                speaker_embed: Optional[Tensor]=None,
                style_embed: Optional[Tensor]=None):
        return self.inference(text_inputs, speaker_embed, style_embed)
</code></pre>
<p>And calling optimize_for_mobile without the preserved_methods parameter:</p>
<pre><code class="lang-auto">ts_prosody_mobile = optimize_for_mobile(ts_prosody)
</code></pre>
<p>ts_prosody_mobile.forward() can be called.<br>
Exporting with torch.jit.save and then loading into Java again:</p>
<pre><code class="lang-auto">        IValue[] prosodyOutput = modelProsody.forward(
                IValue.from(textInputTensor),
                IValue.from(speakerInputTensor),
                IValue.from(styleInputTensor)).toTuple();
</code></pre>
<p>gives me</p>
<pre><code class="lang-auto">com.facebook.jni.CppException: Method 'forward' is not defined.
    Exception raised from get_method at ../../../../src/main/cpp/libtorch_include/arm64-v8a/torch/csrc/jit/api/object.h:103 (most recent call first):
    (no backtrace available)
</code></pre>
<p>Again the same code loading the model without optimize_for_mobile works.</p>
<p>So I ended up instead refactoring my classes from forward/inference to forward_train/forward and now it works.</p>
<p>Only that (without any further quantization or fuse calls) the optimize_for_mobile version runs 5 seconds for a given task while the unoptimized one runs for 4 seconds.</p>

username: kimishpatel
cooked: <aside class="quote no-group" data-username="m-toman" data-post="3" data-topic="129033">
<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/m-toman/48/41136_2.png" class="avatar"> m-toman:</div>
<blockquote>
<p>the optimize_for_mobile version runs 5 seconds for a given task while the unoptimized one runs for 4 seconds.</p>
</blockquote>
</aside>
<p>Without knowing the details of the model it might be a bit harder. Are you profiling only one call to forward or multiple?</p>

username: m-toman
cooked: <p>Multiple calls but it’s 3 different models in sequence that are called.</p>
<p>Model A and B are rather heavy on GRUs and LSTMs together with a few (1D) convolutions.<br>
B also uses some Gaussian upsampling mechanism with a bit of manual implementation.<br>
Both are significantly slower in the “optimize for mobile” version.<br>
The biggest part of it is this layer</p><aside class="onebox githubblob" data-onebox-src="https://github.com/as-ideas/ForwardTacotron/blob/203daff3e90536d8c4f8ca8f28df5364361ce1d1/models/common_layers.py#L39">
  <header class="source">

      <a href="https://github.com/as-ideas/ForwardTacotron/blob/203daff3e90536d8c4f8ca8f28df5364361ce1d1/models/common_layers.py#L39" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/as-ideas/ForwardTacotron/blob/203daff3e90536d8c4f8ca8f28df5364361ce1d1/models/common_layers.py#L39" target="_blank" rel="noopener nofollow ugc">as-ideas/ForwardTacotron/blob/203daff3e90536d8c4f8ca8f28df5364361ce1d1/models/common_layers.py#L39</a></h4>



  <pre class="onebox">    <code class="lang-py">
      <ol class="start lines" start="29" style="counter-reset: li-counter 28 ;">
          <li>        self.conv = nn.Conv1d(in_channels, out_channels, kernel, stride=1, padding=kernel // 2, bias=False)</li>
          <li>        self.bnorm = nn.BatchNorm1d(out_channels)</li>
          <li>        self.relu = relu</li>
          <li>
          <li>    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:</li>
          <li>        x = self.conv(x)</li>
          <li>        x = F.relu(x) if self.relu is True else x</li>
          <li>        return self.bnorm(x)</li>
          <li>
          <li>
          <li class="selected">class CBHG(nn.Module):</li>
          <li>
          <li>    def __init__(self,</li>
          <li>                 K: int,</li>
          <li>                 in_channels: int,</li>
          <li>                 channels: int,</li>
          <li>                 proj_channels: list,</li>
          <li>                 num_highways: int,</li>
          <li>                 dropout: float = 0.5) -&gt; None:</li>
          <li>        super().__init__()</li>
          <li>
      </ol>
    </code>
  </pre>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<p>
The rest is also quite similar to the model in this repo.</p>
<p>Model C is mostly transposed 1D convolutions and here the performance is pretty much the same (looks like this <a href="https://github.com/descriptinc/melgan-neurips/blob/6488045bfba1975602288de07a58570c7b4d66ea/mel2wav/modules.py#L104" class="inline-onebox" rel="noopener nofollow ugc">melgan-neurips/modules.py at 6488045bfba1975602288de07a58570c7b4d66ea · descriptinc/melgan-neurips · GitHub</a>)</p>