[pytorch] torch.nn.Module.parameters(), named_parameters()

오늘은 torch.nn.Module.parameters(), named_parameters()를 이용하여 Deep Neural Network의 convolution layer의 parameters인 weight, bias를 추출해보겠습니다.

 

오늘의 내용은 아래와 같습니다.

< INDEX >

1. torch.nn.Module.parameters() 설명

 

2. torch.nn.Module.named_parameters() 설명

---

제가 구성한 model은 아래 코드와 같습니다.

- 구현 코드 -

import torch
import torch.nn as nn

class CNN(nn.Module):
    def __init__(self, img_size, num_class):
        super(CNN, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 32, 3, 1, 1),
            nn.LeakyReLU(0.2),
            nn.Conv2d(32, 64, 3, 1, 1),
            nn.LeakyReLU(0.2),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(64, 128, 3, 1, 1),
            nn.LeakyReLU(0.2),
            nn.Conv2d(128, 256, 3, 1, 1),
            nn.LeakyReLU(0.2),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.LeakyReLU(0.2),
            nn.MaxPool2d(2,2),
            nn.Conv2d(512, num_class, 3, 1, 1),
            nn.LeakyReLU(0.2)
        )
        
        self.avg_pool = nn.AvgPool2d(img_size // 8)
        self.classifier = nn.Linear(num_class, num_class)

    def forward(self, x):
        features = self.conv(x)
        flatten = self.avg_pool(features).view(features.size(0), -1)
        output = self.classifier(flatten)
        return output, features

1. torch.nn.Module.parameters()

- nn.Module.parameters()는 명칭대로 convolution layer의 parameters value를 추출하는 함수입니다.

 

- 위의 모델에서 존재하는 conv layer는 nn.Conv2d()와 nn.Linear() 입니다. 따라서, nn.Sequential() 내부에 존재하는 6개의 conv layer와 fully connected layer인 Linear layer의 weight 및 bias의 value를 출력할 수 있습니다.

 

- 실제로 Model layer의 구성을 확인하면 아래와 같습니다.

- 여기서 'conv', 'classifier' name으로 정의된 layer의 parameters를 뽑는 것입니다.

 

- 따라서, 각 params의 출력 shape을 확인해보면, 아래와 같이 확인됩니다.

 

- 결과적으로 각 params[idx]를 통해서 각 conv layer에 접근하면, 접근한 conv layer의 parameters(weight, bias)를 추출 가능합니다. 당연히 weight와 bias float 형식의 데이터로 정해질것입니다.

---

2. torch.nn.Module.named_parameters()

- named_parameters()는 제네레이터로 이루어져 각 conv layer의 명칭과 parameters를 동시에 추출할 수 있습니다.