What is the difference between register_parameter and register_buffer in PyTorch?

前端 未结 2 2028
感动是毒
感动是毒 2020-12-30 06:12

Module\'s parameters get changed during training, that is, they are what is learnt during training of a neural network, but what is a buffer?

and is it learnt during

2条回答
  •  猫巷女王i
    2020-12-30 07:06

    Both parameters and buffers you create for a module (nn.Module).

    Say you have a linear layer nn.Linear. You already have weight and bias parameters. But if you need a new parameter you use register_parameter() to register a new named parameter that is a tensor.

    When you register a new parameter it will appear inside module.parameters() iterator, but when you register a buffer it will not.

    The difference:

    Buffers are named tensors that do not update gradients at every step, like parameters. For buffers, you create your custom logic (fully up to you).

    The good thing is when you save the model, all params and buffers are saved, and when you move the model to or off the CUDA params and buffers will go as well.

提交回复
热议问题