Gradients are all zero

Hi,

I was trying to write a custom callback to synthesize sound from gradients during training as in this recent post by Christian Perone: Listening to the neural network gradient norms during training. Here’s what I have so far:

@dataclass
class GenSound(Callback):
    def __init__(self, learn:Learner, fs:int=44100, duration:float=0.01, f:float = 200.0, n_iterations:int=200, fname:str='nn_sound.wav')->None:
        super().__init__()
        self.learn = learn
        self.fs = fs
        self.duration = duration
        self.f = f
        self.n_iterations=n_iterations
        self.fname = fname
        self.p, self.stream = self.open_stream()
        self.frames = []
                
    def on_step_end(self, **kwargs:Any)->None:
        norms = []
        for layer in self.learn.model.ordered_layers:
            norm_grad = layer.weight.grad.norm()
            norms.append(norm_grad)
            print(norm_grad)

            tone = self.f + ((norm_grad.cpu().numpy()) * 100.0)
            tone = tone.astype(np.float32)
            samples = self.generate_tone(tone)

            self.frames.append(samples)

        silence = np.zeros(samples.shape[0] * 2,
                           dtype=np.float32)
        self.frames.append(silence)
        
        if kwargs['iteration'] == 200:
            wf = wave.open(self.fname, 'wb')
            wf.setnchannels(1)
            wf.setsampwidth(self.p.get_sample_size(pyaudio.paFloat32))
            wf.setframerate(self.fs)
            wf.writeframes(b''.join(self.frames))
            wf.close()

            self.stream.stop_stream()
            self.stream.close()
            self.p.terminate()
            
            stop_epoch=True
            stop_training=True
            return {'stop_epoch':stop_epoch, 'stop_training':stop_training}
    
        return

    def on_train_end(self, **kwargs:Any)->None:
        pass
        
    def open_stream(self):
        p = pyaudio.PyAudio()
        stream = p.open(format=pyaudio.paFloat32,
                        channels=1,
                        rate=self.fs,
                        output=True)
        return p, stream


    def generate_tone(self, freq):
        npsin = np.sin(2 * np.pi * np.arange(self.fs*self.duration) * freq / self.fs)
        samples = npsin.astype(np.float32)
        return 0.1 * samples

learn.model.ordered_layers is:

[
            self.features.denseblock1.denselayer6.conv2,
            self.features.denseblock2.denselayer12.conv2,
            self.features.denseblock3.denselayer24.conv2,
            self.features.denseblock4.denselayer16.conv2
]

However, the gradients (norm_grad) assigned in the function on_step_end is always zero.

What is the correct way to access the gradients in a callback?

Thanks :slight_smile: