Python PyTorch ScriptModule.register_buffer用法及代碼示例

本文簡要介紹python語言中 torch.jit.ScriptModule.register_buffer 的用法。

用法:

register_buffer(name, tensor, persistent=True)

參數：

• name(string) -緩衝區的名稱。可以使用給定名稱從此模塊訪問緩衝區

• tensor(Tensor或者None) -要注冊的緩衝區。如果None，然後在緩衝區上運行的操作，例如cuda, 被忽略。如果None，緩衝區是不是包含在模塊的state_dict.

• persistent(bool) -緩衝區是否是該模塊state_dict的一部分。

向模塊添加緩衝區。

這通常用於注冊不應被視為模型參數的緩衝區。例如，BatchNorm 的 running_mean 不是參數，而是模塊狀態的一部分。默認情況下，緩衝區是持久的，並且將與參數一起保存。可以通過將 persistent 設置為 False 來更改此行為。持久緩衝區和非持久緩衝區之間的唯一區別是後者不會成為該模塊 state_dict 的一部分。

可以使用給定名稱將緩衝區作為屬性訪問。

例子：

>>> self.register_buffer('running_mean', torch.zeros(num_features))