PyTorch nn.Module中的self.register_buffer()解析

PyTorch中定义模型时,有时候会遇到self.register_buffer('name', Tensor)的操作,该方法的作用是定义一组参数,该组参数的特别之处在于:模型训练时不会更新(即调用 optimizer.step() 后该组参数不会变化,只可人为地改变它们的值),但是保存模型时,该组参数又作为模型参数不可或缺的一部分被保存。

为了更好地理解这句话,按照惯例,我们通过一个例子实验来解释:

首先,定义一个模型并实例化:

import torch 
import torch.nn as nn
from collections import OrderedDict

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        # (1)常见定义模型时的操作
        self.param_nn = nn.Sequential(OrderedDict([
            ('conv', nn.Conv2d(1, 1, 3, bias=False)),
            ('fc', nn.Linear(1, 2, bias=False))
        ]))

        # (2)使用register_buffer()定义一组参数
        self.register_buffer('param_buf', torch.randn(1, 2))

        # (3)使用形式类似的register_parameter()定义一组参数
        self.register_parameter('param_reg', nn.Parameter(torch.randn(1, 2)))

        # (4)按照类的属性形式定义一组变量
        self.param_attr = torch.randn(1, 2) 

    def forward(self, x):
        return x

net = Model()

上例中,我们通过继承nn.Module类定义了一个模型,在模型参数的定义中,我们分别以(1)常见的nn.Module类形式、(2)self.register_buffer()形式、(3)self.register_parameter()形式,以及(4)python类的属性形式定义了4组参数。

(1)哪些参数可以在模型训练时被更新?

这可以通过net.parameters()查看,因为定义优化器时是这样的:optimizer = SGD(net.parameters(), lr=0.1)。为了方便查看,我们使用 net.named_parameters():

In [8]: list(net.named_parameters())
Out[8]:
[('param_reg',
  Parameter containing:
  tensor([[-0.0617, -0.8984]], requires_grad=True)),
 ('param_nn.conv.weight',
  Parameter containing:
  tensor([[[[-0.3183, -0.0426, -0.2984],
            [-0.1451,  0.2686,  0.0556],
            [-0.3155,  0.0451,  0.0702]]]], requires_grad=True)),
 ('param_nn.fc.weight',
  Parameter containing:
  tensor([[-0.4647],
          [ 0.7753]], requires_grad=True))]

可以看到,我们定义的4组参数中,只有(1)和(3)定义的参数可以被更新,而self.register_buffer()和以python类的属性形式定义的参数都不能被更新。也就是说,modules和parameters可以被更新,而buffers和普通类属性不行。

那既然这两种形式定义的参数都不能被更新,二者可以互相替代吗?答案是不可以,原因看下一节:

(2)这其中哪些才算是模型的参数呢?

模型的所有参数都装在 state_dict 中,因为保存模型参数时直接保存 net.state_dict()。我们看一下其中究竟是哪些参数:

In [9]: net.state_dict()
Out[9]:
OrderedDict([('param_reg', tensor([[-0.0617, -0.8984]])),
             ('param_buf', tensor([[-1.0517,  0.7663]])),
             ('param_nn.conv.weight',
              tensor([[[[-0.3183, -0.0426, -0.2984],
                        [-0.1451,  0.2686,  0.0556],
                        [-0.3155,  0.0451,  0.0702]]]])),
             ('param_nn.fc.weight',
              tensor([[-0.4647],
                      [ 0.7753]]))])

可以看到,通过 nn.Module 类、self.register_buffer() 以及 self.register_parameter() 定义的参数都在 state-dict 中,只有用python类的属性形式定义的参数不包含其中。也就是说,保存模型时,buffers,modules和parameters都可以被保存,但普通属性不行。

(3)self.register_buffer() 的使用方法

在用self.register_buffer('name', tensor) 定义模型参数时,其有两个形参需要传入。第一个是字符串,表示这组参数的名字;第二个就是tensor 形式的参数。

在模型定义中调用这个参数时(比如改变这组参数的值),可以使用self.name 获取。本文例中,就可用self.param_buf 引用。这和类属性的引用方法是一样的。

在实例化模型后,获取这组参数的值时,可以用 net.buffers() 方法获取,该方法返回一个生成器(可迭代变量):

In [10]: net.buffers()
Out[10]: <generator object Module.buffers at 0x00000289CA0032E0>

In [11]: list(net.buffers())
Out[11]: [tensor([[-1.0517,  0.7663]])]

# 也可以用named_buffers() 方法同时获取名字
In [12]: list(net.named_buffers())
Out[12]: [('param_buf', tensor([[-1.0517,  0.7663]]))]

(4)modules, parameters 和 buffers

实际上,PyTorch 定义的模型用OrderedDict() 的方式记录这三种类型,分别保存在self._modules, self._parameters 和 self._buffers 三个私有属性中。调试模式时就可以看到每个模型都有这几个私有属性:


调试模式 变量窗口

由于是私有属性,我们无法在实例化的变量上调用这些属性,可以在模型定义中调用它们:

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        # 常见定义模型时的操作
        self.param_nn = nn.Sequential(OrderedDict([
            ('conv', nn.Conv2d(1, 1, 3, bias=False)),
            ('fc', nn.Linear(1, 2, bias=False))
        ]))

        # 使用register_buffer()定义一组参数
        self.register_buffer('param_buf', torch.randn(1, 2))

        # 使用形式类似的register_parameter()定义一组参数
        self.register_parameter('param_reg', nn.Parameter(torch.randn(1, 2)))

        # 按照类的属性形式定义一组变量
        self.param_attr = torch.randn(1, 2) 

        print('self._modules: ', self._modules)
        print('self._parameters: ', self._modules)
        print('self._buffers: ', self._modules)

    def forward(self, x):
        return x

模型实例化时,调用了 init() 方法,我们就可以看到调用输出结果:

In [21]: net = Model()
self._modules:  OrderedDict([('param_nn', Sequential(
  (conv): Conv2d(1, 1, kernel_size=(3, 3), stride=(1, 1), bias=False)
  (fc): Linear(in_features=1, out_features=2, bias=False)
))])

self._parameters:  OrderedDict([('param_reg', Parameter containing:
tensor([[-0.5666, -0.2624]], requires_grad=True))])

self._buffers:  OrderedDict([('param_buf', tensor([[-0.4005, -0.8199]]))])

在模型的实例化变量上调用时,三者有着相似的方法:

net.modules()
net.named_modules()

net.parameters()
net.named_parameters()

net.buffers()
net.named_buffers()

细心的读着可能会发现,self._parameters 和 net.parameters() 的返回值并不相同。这里self._parameters 只记录了使用 self.register_parameter() 定义的参数,而net.parameters() 返回所有可学习参数,包括self._modules 中的参数和self._parameters 参数的并集。

实际上,由nn.Module类定义的参数和self.register_parameter() 定义的参数性质是一样的,都是nn.Parameter 类型。

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 204,732评论 6 478
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 87,496评论 2 381
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 151,264评论 0 338
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 54,807评论 1 277
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 63,806评论 5 368
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 48,675评论 1 281
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 38,029评论 3 399
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 36,683评论 0 258
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 41,704评论 1 299
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,666评论 2 321
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,773评论 1 332
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,413评论 4 321
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 39,016评论 3 307
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 29,978评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,204评论 1 260
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 45,083评论 2 350
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 42,503评论 2 343

推荐阅读更多精彩内容