pytorch的register_buffer

Version pytorch 1.7
Belong torch.nn.Module


Adds a buffer to the module.
This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.
Buffers can be accessed as attributes using given names.

有时候会面临这个一个情况,有些变量需要伴随着module,但是它又不是module的一个参数,也就是这个参数不需要迭代更新。比如BN的running_mean 是一个伴随着module的必要参数,但是这个参数不需要梯度下降进行更新。

register_buffer有两种用法,如果此参数需要添加到state_dict则需要设置为persistent否则不需要

你可能感兴趣的:(pytorch)