Parameters stuck at zero for custom layer

rselagam: <p>Hello, I’m a beginner at using pytorch, so I’m not sure if I’m approaching this correctly. I have a layer at the beginning of my network to apply an acoustic filter on the magnitude of a FFT input of an audio signal (see equation C.171 from <a href="https://ccrma.stanford.edu/~jos/pasp/Wave_Impedance_Cone.html" rel="noopener nofollow ugc">here</a> for more details).</p>
<p>Here’s what I have for my layer implementation:</p>
<pre><code class="lang-auto">class ConeFilter(torch.nn.Module):
    def __init__(self, 
        n_fft=512, 
        sample_rate=44100.0,
        steps=100):
        super().__init__()
        self.x0 = torch.nn.Parameter(torch.zeros(1))
        self.angle = torch.nn.Parameter(torch.zeros(1))
        self.depth = torch.nn.Parameter(torch.zeros(1))
        self.freq_map = torch.from_numpy(librosa.fft_frequencies(sr=sample_rate, n_fft=n_fft)).float()
        self.c = 343 # Speed of sound in air m/s
        self.rho = 1.293 # Air density kg/m^3
        self.steps = 100
        self.pi = 3.1415927

    def forward(self, noisy):
        scalar = self.c*self.rho
        result = noisy
        self.freq_map = self.freq_map.to('cuda:0')
        tot_impedance = None
        myx0 = torch.relu(self.x0) # x0 must be non-negative
        mydepth = torch.relu(self.depth) # depth must be non-negative
        myangle = ((torch.sigmoid(self.angle) + 1.0)/2.0) * (self.pi/2.0)

        for i in range(self.steps):
            x = (myx0 + (i * mydepth/self.steps)).float()
            r = torch.tan(myangle) * x
            numer = torch.mul(self.freq_map, x)
            denom = torch.add(numer, self.c)
            frac = torch.div(numer,denom)
            
            impedance = frac * self.rho * self.c / (self.pi * r * r)         

            if (tot_impedance == None):
                tot_impedance = impedance
            else:
                tot_impedance = torch.mul(tot_impedance, impedance)
        tot_impedance = tot_impedance.unsqueeze(1)
        tot_impedance = tot_impedance.unsqueeze(0)
        result = torch.mul(result, tot_impedance)
        self.print_parameters()
        return result

    def print_parameters(self):
        print("x0: " + str(self.x0.data))
        print("angle: " + str(self.angle.data * self.pi / 180.0))
        print("depth: " + str(self.depth.data))
</code></pre>
<p>The idea is that I want the model to converge on the optimal cone that is expressed by the three terms: <code>x0</code>, <code>angle</code>, and <code>depth</code>. My problem is that the output of the <code>print_parameters()</code> function shows that the data values are always stuck at <code>0.0</code>. I read on another post that accessing the <code>.data</code> field may break the autograd computation graph, but I’m not sure how else to verify that the weights are changing appropriately through training. Any insight would be greatly appreciated.</p>

ariG23498: <p>Just a thought:<br>
Would it be better to use <code>self.x0.detach().item()</code> instead of <code>self.x0.data</code> here? My intuition is that when we would be printing the parameters we would first detach the parameters from gradient flow and then access the item to print. This would not interfere with the gradients.</p>

rselagam: <p>Hmm I changed the calls in the <code>print_parameters</code> function as you described. Values still seem to be stuck at 0. I was concerned about looping within a forward pass, and if that would affect the computation graph. Would this be a problem?</p>

ariG23498: <p>Interesting. I don’t think looping should create a problem. Can you initialize the parameters to something other than 0 and check?</p>

ptrblck: <p>I would also suggest to initialize parameters with random values as the zero init could cause the training to fail.</p>

rselagam: <p>Oh, is that something I would do when I instantiate the parameter in the init class? Do you have a pointer to documentation/could you provide a short code snippet on how to do that?</p>

rselagam: <p>Also, as an aside, although the ConeFilter is a layer created in my networks <code>torch.nn.ModuleList</code>, when I send the net to the GPU with the following:</p>
<pre><code class="lang-auto">net = Network(
                    args.threshold,
                    args.tau_grad,
                    args.scale_grad,
                    args.dmax,
                    args.out_delay)
        net = torch.nn.DataParallel(net.to(device), device_ids=args.gpu)
</code></pre>
<p>The <code>freq_map</code> remains on the CPU, so I have to manually send it to the GPU at the beginning of every forward pass (<code>self.freq_map = self.freq_map.to('cuda:0')</code>). Is there any way around this?</p>

ptrblck: <aside class="quote no-group" data-username="rselagam" data-post="6" data-topic="201893">
<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/rselagam/48/75147_2.png" class="avatar"> rselagam:</div>
<blockquote>
<p>Do you have a pointer to documentation/could you provide a short code snippet on how to do that?</p>
</blockquote>
</aside>
<p>You could use any factory method, e.g. <code>torch.randn</code>, or you could also initialize it via <code>torch.nn.init</code> methods.</p>
<aside class="quote no-group" data-username="rselagam" data-post="7" data-topic="201893">
<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/rselagam/48/75147_2.png" class="avatar"> rselagam:</div>
<blockquote>
<p>The <code>freq_map</code> remains on the CPU</p>
</blockquote>
</aside>
<p><code>self.freq_map</code> is a plain tensor and thus not registered to the module.<br>
I assume you don’t want to train this tensor, so register it as a buffer via <code>self.register_buffer</code> and the <code>to()</code> calls will move it to the GPU as well.</p>

rselagam: <p>Hm, so I used <code>torch.randn(1)</code> in the <code>Parameter</code> instantiation statements, and now I can see that the initial values of the parameters are no longer 0.0. But, the problem remains that the values of the parameters are stuck and don’t change as training progresses. I’m using the <code>RAdam</code> optimizer module. Below is a simplified version of the code I’m running in each training iteration.</p>
<pre><code class="lang-auto">denoised_abs = net(noisy_abs)
clean_rec = stft_mixer(denoised_abs, noisy_arg, args.n_fft)
score = si_snr(clean_rec, clean)
loss = lam * F.mse_loss(denoised_abs, clean_abs) + (100 - torch.mean(score))
if torch.isnan(loss).any():
                loss[torch.isnan(loss)] = 0
assert torch.isnan(loss) == False

optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(net.parameters(), args.clip)
optimizer.step()
</code></pre>
<p>P.S. Thanks for the tip with <code>register_buffer</code>, that part seems to be working better now!</p>