SpeechBrain库grad为NoneType的问题

The SpeechBrain Toolkit

SpeechBrain是一个基于PyTorch的开源语音库,致力于提供一个简单、灵活、用户友好的工具包,可用于开发state-of-the-art的语音系统,包括语音内容识别、说话人识别等。

可参考:GitHub – The SpeechBrain Toolkit

grad为NoneType的问题

SpeechBrain的模型参数是不会有这个问题的,毕竟人家也是这么训练过来的,梯度为None的话还怎么训练。但是对输入音频求梯度却是None,问题演示如下:

考虑对输入求梯度,而不是对模型参数求梯度主要是为了做一些对抗样本(Adversarial Example)的研究。
如果使用Pretrained模型,首先可能想到的是把freeze_params参数设成False,如下:

具体可以参考speechbrain/pretrained/interface.py文件(GitHub – interface.py)。EncoderClassifier类继承了Pretrained类。Pretrained类中freeze_params用于将所有参数的requires_grad属性设为False,如下面的源码所示:

但这其实不是问题所在。熟悉PyTorch的Autograd机制就会知道,如果叶子结点xrequires_grad属性设为Trueyrequires_grad属性设为Falsez=x+yrequires_grad属性会自动为True同时增加一个grad_fn属性,可以求zx的梯度x.grad。注意到这个过程是不需要yrequires_grad也设为True的。



计算图

同样的道理,z=model(x)要计算x的梯度当然也不需要model中的参数设置requires_grad。于是思路就是好好研究一下EncoderClassifier类的几个方法,究竟是哪一步导致grad没有了。EncoderClassifier类的结构如下:

总的逻辑过程在classify_batch(self, wavs, wav_lens=None)classify_file(self, path)中。这两个方法的区别是,前者的输入是一个准备好了的audio向量,而后者的输入是一个音频文件的路径。在craft对抗样本的时候主要使用classify_batch()

第5行用前端(Front-end)编码器(Xvector之类的)将输入wavs编码成embeddings,第6行将emb输入后端(Back-end)分类器,得到各个label的概率分布。第7,8行得到判别结果。我们需要的是out_probwavs的梯度,但是结果就是None

那么问题可以这么来排除,可以将emb设为叶子结点,看梯度能否从out_prob传播回到emb,如果可以,那么就可以排除self.modules.classifier(emb).squeeze(1)的问题。

emb.grad是有的,说明self.modules.classifier(emb).squeeze(1)没有问题。那么问题只能是在self.encode_batch(wavs, wav_lens)了,看看它的源码:

主要问题集中在这8,9,10三行。利用上面同样的问题排除方法,一个一个函数排除,最后可以发现问题出在self.modules.compute_features(wavs)这个方法中。变量只要进过这个方法,梯度就再也无法反向传播了。这个方法是通过hparams_file参数在.yaml文件中定义的。于是去打开hparams_inference.yaml一看,发现这个方法的定义如下:

查阅SpeechBrain的官方文档,可以找到speechbrain.lobes.features.Fbank的源码,可参考 SpeechBrain – Source code for speechbrain.lobes.features

问题所在

于是就能找到问题所在:震惊!居然有个torch.no_grad()。这能梯度传播就有鬼了呀。

以下摘自PyTorch文档:

参考:PyTorch – NO_GRAD

  • Context-manager that disabled gradient calculation.

  • Disabling gradient calculation is useful for inference, when you are sure that you will not call Tensor.backward(). It will reduce memory consumption for computations that would otherwise have requires_grad=True.

  • In this mode, the result of every computation will have requires_grad=False, even when the inputs have requires_grad=True.

到这一步,实际上就能根据关键词找到同道中人了,也有人遇到了同样的问题,并且他跟我的分析非常相似Tracking gradient w.r.t. to input audio sample: Fbank breaks computation graph。虽然他也没有很优雅的解决办法。但至少可以感觉到问题的就在这里。

问题找到,接下来就是解决问题。第一个想法,SpeechBrain有没有能求grad的FilterBank函数?查一下文档还真有,叫speechbrain.processing.features.Filterbank。但是仔细看看的话就会发现,speechbrain.lobes.features.Fbank实际上用到了它。

意思就是,只把speechbrain.lobes.features.Fbank换成speechbrain.processing.features.Filterbank是不可行的,因为speechbrain.lobes.features.Fbank里面还做了一些其他的操作。那最后的办法就是——自力更生,改写这个类的forward方法。修改后的代码如下:

这个类在一般情况下与原来一模一样,只有在实例化的时候将参数allow_grad显式地设为True,这个类就可以支持梯度反向传播。

接下来就是在.yaml文件中定义新的compute_features,如下:

注意别忘了allow_grad

Done!

Leave a Comment