torchvisionから学習済みモデルを使用したい場合に、
・任意の層まで、パラメータを固定したい
・学習済み重みののっかった層をとってきたい
なんてことがありました。
前者はファインチューニングの際に必要ですし、後者はSkip Connection的な実装をする際に必要となってきます。
（Encoderで学習済みモデルを利用して、Upsampling時にEncoderの情報を使うようなシーン等）
というわけで、これらのやり方をメモしておこうと思います。
任意の層までパラメータを固定したい
まずは、サンプルとして vgg16 を持ってきます。
|
1 2 3 |
import torchvision vgg = torchvision.models.vgg16(pretrained=True) |
モデルのパラメータは vgg.parameters() で取り出すことができます。
Generatorオブジェクトですので、いったんリストに格納して中身を見てみます。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
params = list(vgg.parameters()) print(len(params)) ``` 32 ``` for p in params: print(p.size()) ``` torch.Size([64, 3, 3, 3]) torch.Size([64]) torch.Size([64, 64, 3, 3]) torch.Size([64]) torch.Size([128, 64, 3, 3]) torch.Size([128]) torch.Size([128, 128, 3, 3]) torch.Size([128]) torch.Size([256, 128, 3, 3]) torch.Size([256]) torch.Size([256, 256, 3, 3]) torch.Size([256]) torch.Size([256, 256, 3, 3]) torch.Size([256]) torch.Size([512, 256, 3, 3]) torch.Size([512]) torch.Size([512, 512, 3, 3]) torch.Size([512]) torch.Size([512, 512, 3, 3]) torch.Size([512]) torch.Size([512, 512, 3, 3]) torch.Size([512]) torch.Size([512, 512, 3, 3]) torch.Size([512]) torch.Size([512, 512, 3, 3]) torch.Size([512]) torch.Size([4096, 25088]) torch.Size([4096]) torch.Size([4096, 4096]) torch.Size([4096]) torch.Size([1000, 4096]) torch.Size([1000]) ``` |
VGG16なので、要素数は16個かと思いきや、32個あります。
各パラメータのshape を見てあげると、これは重みとバイアスでそれぞれ別にあることがわかります。
各パラメータはデフォルトでrequires_grad=True となっており、再学習されてしまいます。
|
1 2 3 4 5 |
params[0].requires_grad ``` True ``` |
ですので、例えば12層目まで重みバイアスを固定で、それ以降は再学習可能にしたい場合はfor文を回してあげて、任意の層でrequires_grad=Falseとしてあげればよいことになります。
|
1 2 3 |
for i, param in enumerate(vgg.parameters()): if i >= 24: param.requires_grad=False |
学習済み重みののっかった層をとってきたい
例えば、VGGの5層目を通した後の出力を保存しておいて、Decoderのある層に加算しようだとかいう場合。
モデルの重みは、.parameters()で取得することができましたが、定義された層の情報は、.children()で取得することができます。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
layers = list(vgg.children()) print(layers) ``` [Sequential( (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (1): ReLU(inplace=True) (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (3): ReLU(inplace=True) (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (6): ReLU(inplace=True) (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (8): ReLU(inplace=True) (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (11): ReLU(inplace=True) (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (13): ReLU(inplace=True) (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (15): ReLU(inplace=True) (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (18): ReLU(inplace=True) (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (20): ReLU(inplace=True) (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (22): ReLU(inplace=True) (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (25): ReLU(inplace=True) (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (27): ReLU(inplace=True) (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (29): ReLU(inplace=True) (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) ), AdaptiveAvgPool2d(output_size=(7, 7)), Sequential( (0): Linear(in_features=25088, out_features=4096, bias=True) (1): ReLU(inplace=True) (2): Dropout(p=0.5, inplace=False) (3): Linear(in_features=4096, out_features=4096, bias=True) (4): ReLU(inplace=True) (5): Dropout(p=0.5, inplace=False) (6): Linear(in_features=4096, out_features=1000, bias=True) )] ``` |
さらに、中は大きく３つの要素で構成されていることがわかります。
一つ目の、nn.Sequential() の中身はConv層で、特徴抽出器の部分。
二つ目は AdaptiveAvgpool2d()
三つ目は、nn.Sequential()で中身は全結合層なので、分類器の部分ですね。
というわけですので、学習済み重みがセットされた層の、はじめの５層だけを取り出す際には、
|
1 2 3 4 5 6 7 8 9 10 11 |
print(layers[0][:5]) ``` Sequential( (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (1): ReLU(inplace=True) (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (3): ReLU(inplace=True) (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) ) ``` |
こんな風にすればnn.Sequentialのまとまりで層をとってくることができました。
ちなみにtorchvisionのソースを見るとこのようになっており、 ( https://pytorch.org/docs/stable/_modules/torchvision/models/vgg.html )
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
class VGG(nn.Module): def __init__(self, features, num_classes=1000, init_weights=True): super(VGG, self).__init__() self.features = features self.avgpool = nn.AdaptiveAvgPool2d((7, 7)) self.classifier = nn.Sequential( nn.Linear(512 * 7 * 7, 4096), nn.ReLU(True), nn.Dropout(), nn.Linear(4096, 4096), nn.ReLU(True), nn.Dropout(), nn.Linear(4096, num_classes), ) if init_weights: self._initialize_weights() def forward(self, x): x = self.features(x) x = self.avgpool(x) x = torch.flatten(x, 1) x = self.classifier(x) return x def _initialize_weights(self): for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) elif isinstance(m, nn.Linear): nn.init.normal_(m.weight, 0, 0.01) nn.init.constant_(m.bias, 0) |
vggモデルでいえば、同様のことは vgg.features[:5] でも行えますね。
最後まで読んでいただきありがとうございました。
