Reshaping torch.nn.parameter

drj3122: <p>I’m writing a class that is basically just <code>nn.Linear</code> except it has some reshaping.</p>
<pre><code class="lang-auto">class myLinear2(nn.Module):
  def __init__(self, in_features, out_features):
    super().__init__()
    self.in_features = in_features
    self.out_features = out_features
    self.weight = torch.nn.Parameter(torch.randn(out_features, in_features))    
    self.bias = torch.nn.Parameter(torch.randn(out_features))
       
  def forward(self, input):
    self.weight = self.weight.reshape(-1, 1, self.weight.shape[-1])
    self.bias = self.bias.reshape(self.bias.shape[0], 1, 1)      
    if input.dim() &gt; 1:
      x, y = input.shape
      if y != self.in_features:
          sys.exit(f'Wrong Input Features. Please use tensor with {self.in_features} Input Features')
    output = (input @ self.weight.transpose(-1, -2)) + self.bias
    return output
</code></pre>
<p>Running a model with this class returns an error:<br>
<code>TypeError: cannot assign 'torch.cuda.FloatTensor' as parameter 'weight' (torch.nn.Parameter or None expected)</code></p>
<p>However, if I put the reshape in an <code>if</code> statement, it works whether or not the <code>if</code> statement is actually true.</p>
<pre><code class="lang-auto"># this version works
if self.weight.dim() &gt; 2:      
      self.weight = self.weight.reshape(-1, 1, self.weight.shape[-1])
      self.bias = self.bias.reshape(self.bias.shape[0], 1, 1)
</code></pre>
<p>What is going on here?</p>

ptrblck: <p>I don’t think your code would work in the condition as I’m seeing the same error if I increase the number of dimension of <code>self.weight</code> and make sure the code is indeed executed.<br>
I’m not familiar with your use case, but maybe you could avoid re-assigning the reshaped tensor and just use it in the calculation which could avoid this error?<br>
If not, you could try to wrap the assignment into a <code>no_grad()</code> guard and use e.g. <code>self.weight.copy_</code> to avoid creating new <code>nn.Parameters</code> as you would then e.g. need to pass them to an optimizer again.</p>

drj3122: <p>Thanks I’ll try this out. My use case is a little unique since it actually has to do with <code>pyro ppl</code>. I actually have a forum post on their forum <a href="https://forum.pyro.ai/t/reproducible-example-proper-shape-for-predictive-with-parallel-true/4194/6" rel="noopener nofollow ugc">(Reproducible example) Proper shape for Predictive with parallel=True - Pyro Discussion Forum</a>. I was trying to figure out if I could use reshaping to get proper output for the various weight shapes I describe there.</p>

ID56: <p>Let’s simplify the problem:</p>
<pre><code class="lang-python">class Foo(nn.Module):
  def __init__(self):
    super().__init__()
    self.p = nn.Parameter(torch.randn(1))
       
  def forward(self, input):
    self.p = self.p.reshape(1, -1)
    return input * self.p

x = torch.randn(5)
layer = Foo()
y = layer(x)
</code></pre>
<p>This will not work, and your error message hints at why. A simple test we could do:</p>
<pre><code class="lang-auto">def forward(self, input):
    print(type(self.p))
    print(type (self.p.reshape(1,-1)))
    self.p = self.p.reshape(1, -1)
    return input * self.p
</code></pre>
<p>This will show:</p>
<pre><code class="lang-auto">&lt;class 'torch.nn.parameter.Parameter'&gt;
&lt;class 'torch.Tensor'&gt;
</code></pre>
<p>So calling reshape on a Parameter returns the tensor. Then one could think of simply wrapping an nn.Parameter around self.p.reshape(). While that will resolve the error, it would be a new parameter untracked by the optimizer and thus of no use.</p>
<p>One approach to try: we can modify the underlying data tensor.</p>
<pre><code class="lang-auto">def forward(self, input):
    self.p.data = self.p.data.reshape(1, -1)
    return input * self.p
</code></pre>
<p>Or, we could try not modifying the parameter at all, and just storing in a temporary variable.</p>
<pre><code class="lang-auto">def forward(self, input):
    p = self.p.reshape(1, -1)
    return input * p
</code></pre>
<p>Of course, it would be the best if no modification was done on self.p inside forward at all, and the right shape was assigned in <code>__init__</code>.</p>