Parameter order changes when using TorchScript

hungry: <p>Hello,</p>
<p>After updating PyTorch 1.0.1 -&gt; 1.4.0 I encountered some unwanted behavior. When creating a ScriptModule, the order of its Parameters changes. I’ll attach an example based on the <a href="https://github.com/pytorch/benchmark/blob/master/rnns/fastrnns/custom_lstms.py" rel="nofollow noopener">https://github.com/pytorch/benchmark/blob/master/rnns/fastrnns/custom_lstms.py</a> file.</p>
<p>Two classes are defined. The only difference between the two is that one is a ScriptModule and the other an nn.Module. By plotting their parameter sizes we see that they are now stored in a different order than declared in the class (they also stop matching the native cell version). It might cause some backward compatibility problems (e.g. in the case of ‘custom_lstms.py’ the ‘test_script…’ functions stopped working).</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
from torch.nn import Parameter
import torch.jit as jit
from typing import List, Tuple
from torch import Tensor
</code></pre>
<pre><code class="lang-auto">class LSTMCell(jit.ScriptModule):
    def __init__(self, input_size, hidden_size):
        super(LSTMCell, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.weight_ih = Parameter(torch.randn(4 * hidden_size, input_size))
        self.weight_hh = Parameter(torch.randn(4 * hidden_size, hidden_size))
        self.bias_ih = Parameter(torch.randn(4 * hidden_size))
        self.bias_hh = Parameter(torch.randn(4 * hidden_size))

    @jit.script_method
    def forward(self, input, state):
        # type: (Tensor, Tuple[Tensor, Tensor]) -&gt; Tensor
        return input
    

class LSTMCellNoJit(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(LSTMCellNoJit, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.weight_ih = Parameter(torch.randn(4 * hidden_size, input_size))
        self.weight_hh = Parameter(torch.randn(4 * hidden_size, hidden_size))
        self.bias_ih = Parameter(torch.randn(4 * hidden_size))
        self.bias_hh = Parameter(torch.randn(4 * hidden_size))

    def forward(self, input, state):
        return input
</code></pre>
<pre><code class="lang-auto">net = LSTMCell(2, 10)

pars = [p for p in net.parameters()]
for i in range(len(pars)):
    print(pars[i].shape)
</code></pre>
<p>Output:<br>
torch.Size([40, 10])<br>
torch.Size([40])<br>
torch.Size([40])<br>
torch.Size([40, 2])</p>
<pre><code class="lang-auto">net = LSTMCellNoJit(2, 10)

pars = [p for p in net.parameters()]
for i in range(len(pars)):
    print(pars[i].shape)
</code></pre>
<p>Output:<br>
torch.Size([40, 2])<br>
torch.Size([40, 10])<br>
torch.Size([40])<br>
torch.Size([40])</p>

driazati: <p>Hey thanks for the report / repro! This is an issue that’s been open for a while, we’re planning to fix it soon</p>
<aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/issues/25047" target="_blank">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/25047" target="_blank">ScriptModule and nn.Module parameter ordering difference</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2019-08-22" data-time="21:27:58" data-timezone="UTC">09:27PM - 22 Aug 19 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/wanchaol" target="_blank">
          <img alt="wanchaol" src="https://avatars2.githubusercontent.com/u/9443650?v=4" class="onebox-avatar-inline" width="20" height="20">
          wanchaol
        </a>
      </div>
    </div>
  </div>
</div>

<div class="github-row">
  <p class="github-content">class eagerNet(nn.Module):
 def __init__(self):
 super(eagerNet, self).__init__()
 self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
 self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
 self.fc1 = nn.Linear(320, 50)
 self.fc2...</p>
</div>

<div class="labels">
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">jit</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">jit-backlog</span>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>