Pytorch学習(19)---モデルでのbufferの使用

3385 ワード

いつも言う


pytorchは一般的にネットワーク内のパラメータをOrderedDict形式に保存することを知っています.ここで額パラメータには実は2種類が含まれています.1つは、モデルに含まれる様々なmoduleのパラメータ、すなわちnn.Parameterであり、もちろん、ネットワーク内で他のnn.Parameterパラメータを定義することができる.もう一つはbufferです.前者はoptim.stepのたびに更新され、後者は更新されない.

class myModel(nn.Module):
    def __init__(self, kernel_size=3):
        super(Depth_guided1, self).__init__()
        self.kernel_size = kernel_size
        self.back_end = torch.nn.Sequential(
            torch.nn.Conv2d(3, 32, 3, padding=1),
            torch.nn.ReLU(True),
            torch.nn.Conv2d(3, 64, 3, padding=1),
            torch.nn.ReLU(True),
            torch.nn.Conv2d(64, 3, 3, padding=1),
            torch.nn.ReLU(True),
        )

        mybuffer = np.arange(1,10,1)
        self.mybuffer_tmp = np.randn((len(mybuffer), 1, 1, 10), dtype='float32')
        self.mybuffer_tmp = torch.from_numpy(self.mybuffer_tmp)
        # register preset variables as buffer
        # So that, in testing , we can use buffer variables.
        self.register_buffer('mybuffer', self.mybuffer_tmp)

        # Learnable weights
        self.conv_weights = nn.Parameter(torch.FloatTensor(64, 10).normal_(mean=0, std=0.01))
        # Other code
        def forward(self):
            ...
            #   self.mybuffer!

知識点:I.定義Parameterおよびbufferは、いずれもTensorに入力するだけでよい.gpuに変換する必要もありません.これは,ネットワークが.cuda()を行うと,内部の層のパラメータ,bufferなどが自動的に対応するGPUに変換されるためである.
II. self.register_bufferはtensorをbufferとして登録することができ、forwardではself.mybufferを用いるが、self.mybuffer_tmpではない.
III.ネットワーク格納時にはbufferも格納され,ネットワークloadモデルの場合には格納されたモデルのbufferも付与される.
IV.bufferの更新はforwardであり、optim.stepnn.Parameterタイプのパラメータしか更新できない.