スポンサーリンク

【PyTorch】モデルの重みに初期値を設定する方法

Machine Learning

PyTorchモデルの重みに初期値を設定する方法について記載していきます。今回、対象とするモデルはシンプル構造の下記コードとします。

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(3, 3)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.linear(x)
        x = self.relu(x)     
        return x

model = Model()
print(model)
# Model(
#   (linear): Linear(in_features=3, out_features=3, bias=True)
#   (relu): ReLU()
# )

## 重みとバイアスの初期値を確認
for param in model.parameters():
  print(param)

# Parameter containing:
# tensor([[ 0.2242, -0.3569,  0.1490],
#         [-0.2229, -0.1892,  0.3646],
#         [-0.1613, -0.3112,  0.4168]], requires_grad=True)
# Parameter containing:
# tensor([-0.3672, -0.3400, -0.2994], requires_grad=True)

モデルの重みに初期値を設定する

torch.nn.initの利用する方法

torch.nn.initを利用することで、重みの初期値を簡単に設定することができます。

例えば、重みの初期値を全て0にするには、torch.nn.init.zeros_()を用います。また、バイアスの初期値にはtorch.nn.init.ones_()を用いて全て1にします。全ての重みとバイアスが指定した値になっていることが確認できます。

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(3, 3)
        nn.init.zeros_(self.linear.weight) # 重みの初期値を設定
        nn.init.ones_(self.linear.bias)    # バイアスの初期値を設定

        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.linear(x)
        x = self.relu(x)     
        return x

model = Model()
for param in model.parameters():
  print(param)
# Parameter containing:
# tensor([[0., 0., 0.],
#         [0., 0., 0.],
#         [0., 0., 0.]], requires_grad=True)
# Parameter containing:
# tensor([1., 1., 1.], requires_grad=True)

初期値の指定方法には、他にも様々な方法があります。

torch.nn.init.uniform_()は、最小値a、最大値bの範囲の一様分布から初期化します。

torch.nn.init.uniform_(tensor, a=0.0, b=1.0)

torch.nn.init.normal_()は、平均mean、標準偏差stdの正規分布から初期化します。

torch.nn.init.normal_(tensor, mean=0.0, std=1.0)

torch.nn.init.constant_()は、定数valで初期化します。

torch.nn.init.constant_(tensor, val)

Xavier、Heの方法もtorch.nn.init.xavier_normal_()やtorch.nn.init.kaiming_normal_()を用いることで使うことができます。

任意の値を指定する方法

torch.nn.initを使う方法の場合は、一括で重みの値を変えることができますが、任意の値を指定することはできない。任意の値を指定するには、nn.Parameterを用いてTensor配列を定義し、モデルの重みやバイアスに値を代入することで指定できる。

model = Model()

weight_new = nn.Parameter(torch.tensor([[1., 2., 3.],
                                          [4., 5., 6.],
                                          [7., 8., 9.]])) 

bias_new = nn.Parameter(torch.tensor([[10., 11., 12.]])) 

model.linear.weight = weight_new
model.linear.bias = bias_new

for param in model.parameters():
  print(param)
# Parameter containing:
# tensor([[1., 2., 3.],
#         [4., 5., 6.],
#         [7., 8., 9.]], requires_grad=True)
# Parameter containing:
# tensor([[10., 11., 12.]], requires_grad=True)

関連記事、参考資料

コメント