[Pytorch] network(model) parameter 얻는 방법

소개

최근 경량화 스터디를 시작했다. 경량화 스터디에서 첫 주제는 Pruning이었다. Pruning은 가지치기 기법으로 모델에서 중요하지 않은 weight나 filter를 제거함으로써 계산량과 모델 크기를 줄여주는 방법이다. 중요한 weight나 filter를 찾아내기 위해 모델 내부의 parameter에  직접 접근하고, 변경도 해야해서 구현이 쉽지 않았다. 구현을 위해 여러 시도를 해보면서 알게된 점을 정리해보았다.

 

정리

0. state_dict를 사용하여 값을 변경하면 된다. 나의 목표는 parameter 참고와 변경이었으니, get_parameters와 state_dict를 적절히 활용하면 될 것 같다.

 

1. 모듈에 named_parameters / parameters, named_children / children처럼 두 종류가 있는데, named가 붙으면 (name, parameters( or children) )을, 붙지 않으면 parameters( or children) 을 return 한다.

 

2. 확실히 알아보진 않았지만, children을 submodule로 취급하고, nn.으로 시작하는 모듈들을 module취급했다. 그래서 get_children()을 하면 모델의 멤버변수들만을 return 해주었지만, get_modules()를 하면 model 내에 사용된 모든 nn.module 들을 return 해주는 것 같다.

self.children() -> submodule 반환 (모델의 멤버변수)

self.modules() -> nn.modules 반환 ( 모델에 들어간 nn.~모듈)

self.parameters() -> 파라미터 반환 (여기서 파라미터란, weight, bias처럼 torch.tensor로 이루어 진 것들을 의미함.)

 

이후 pruning 코드를 완성하면 포스트하도록 하겠습니다.

cf)

3. class 관련 함수가 다양하게 있었다.  getattr(class, 멤버변수(attr)), hasattr(),delattr() 등등..

코드

사용한 모델인 ToyNet 선언은 아래에 작성해두었다.

 

구글링을 통해서 접근하는 방법 몇가지를 찾을 수 있었다.

1. self.named_parameters()

2. self.children()

3. self.get_parameter(target)

4. state_dict()

5. self.named_modules()

1. self.named_parameters()

for name, param in net.named_parameters():
    print(name,param.shape)

layer1.0.weight torch.Size([16, 3, 3, 3])
layer1.1.weight torch.Size([16])
layer1.1.bias torch.Size([16])
layer2.0.weight torch.Size([32, 16, 3, 3])
layer2.1.weight torch.Size([32])
layer2.1.bias torch.Size([32])
layer3.0.weight torch.Size([64, 32, 3, 3])
layer3.1.weight torch.Size([64])
layer3.1.bias torch.Size([64])
fc.0.weight torch.Size([10, 4096])
fc.0.bias torch.Size([10])

self.named_parameters()는 iteration을 return 해주어 다음과 같이 사용할 수 있다.

2. self.children()

for child in net.children():
    print('===')
    print(child)

===
Sequential(
  (0): Conv2d(3, 3, kernel_size=(3, 3), stride=(1, 1))
  (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (2): ReLU(inplace=True)
)
===
Sequential(
  (0): Conv2d(3, 3, kernel_size=(3, 3), stride=(1, 1))
  (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (2): ReLU(inplace=True)
)
===
Sequential(
  (0): Conv2d(3, 3, kernel_size=(3, 3), stride=(1, 1))
  (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (2): ReLU(inplace=True)
)
===
Sequential(
  (0): Conv2d(3, 3, kernel_size=(3, 3), stride=(1, 1))
)

 

각 child를 return 해준다. 따라서 만약 child 변경이 필요로 하면 이 방법을 사용하면 가능하다.

3. self.get_parameter(target)

net.get_parameter('layer1.0.weight')

Parameter containing:
tensor([[[[ 0.1428, -0.1008,  0.0711],
          [ 0.1136,  0.0253, -0.1493],
          [ 0.0799, -0.1384,  0.1538]],

         [[-0.1913, -0.0300, -0.0979],
          [-0.1769,  0.0280,  0.0517],
          [ 0.1063, -0.1397,  0.0934]],

         [[-0.0517, -0.0134, -0.0593],
          [-0.1694, -0.1802, -0.1040],
          [-0.1396, -0.0950,  0.1722]]],


        [[[-0.0842,  0.0227,  0.0371],
          [-0.1586, -0.1911, -0.0911],
          [-0.0009, -0.1598,  0.0900]],

         [[-0.0630, -0.0168,  0.1889],
          [ 0.0886, -0.0722, -0.1239],
          [-0.0520,  0.0752, -0.1045]],

         [[ 0.1865, -0.0235,  0.0067],
          [ 0.0536, -0.0746, -0.0076],
          [-0.0215,  0.1758, -0.0804]]],


        [[[-0.1117, -0.1731,  0.1504],
          [ 0.1045, -0.0573, -0.0749],
          [-0.1198,  0.0469,  0.1872]],

         [[-0.0679, -0.0966,  0.1902],
          [ 0.0304, -0.0293,  0.1168],
          [-0.0383,  0.1735,  0.0266]],

         [[ 0.1863, -0.1081,  0.1664],
          [-0.0958, -0.0546,  0.0216],
          [ 0.0766, -0.0499, -0.1173]]]], requires_grad=True)

get_parameters()는 target이 모델의 children 중에 존재하면 해당 child의 parameter를 return 해준다.

 

4. self.state_dict()

구글링 도중 state_dict를 발견했다. 

https://pytorch.org/tutorials/beginner/saving_loading_models.html

Saving and Loading Models — PyTorch Tutorials 1.9.1+cu102 documentation

model = ToyNet()

# Initialize optimizer
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# Print model's state_dict
print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

# Print optimizer's state_dict
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])

Model's state_dict:
layer1.0.weight 	 torch.Size([16, 3, 3, 3])
layer1.1.weight 	 torch.Size([16])
layer1.1.bias 	 torch.Size([16])
layer1.1.running_mean 	 torch.Size([16])
layer1.1.running_var 	 torch.Size([16])
layer1.1.num_batches_tracked 	 torch.Size([])
layer2.0.weight 	 torch.Size([32, 16, 3, 3])
layer2.1.weight 	 torch.Size([32])
layer2.1.bias 	 torch.Size([32])
layer2.1.running_mean 	 torch.Size([32])
layer2.1.running_var 	 torch.Size([32])
layer2.1.num_batches_tracked 	 torch.Size([])
layer3.0.weight 	 torch.Size([64, 32, 3, 3])
layer3.1.weight 	 torch.Size([64])
layer3.1.bias 	 torch.Size([64])
layer3.1.running_mean 	 torch.Size([64])
layer3.1.running_var 	 torch.Size([64])
layer3.1.num_batches_tracked 	 torch.Size([])
fc.0.weight 	 torch.Size([10, 4096])
fc.0.bias 	 torch.Size([10])
Optimizer's state_dict:
state 	 {}
param_groups 	 [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]}]

state = net.state_dict()
print(state['layer1.0.weight'].shape)
state['layer1.0.weight']=torch.tensor((3,3))
print(state['layer1.0.weight'].shape)

torch.Size([16, 3, 3, 3])
torch.Size([2])

state_dict는 말그대로 dictionary로, 변경됨을 확인 할 수 있다.

5. self.named_modules() 

for name,layer in net.named_modules():
    print(name,layer)

 ToyNet(
  (layer1): Sequential(
    (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU(inplace=True)
  )
  (layer2): Sequential(
    (0): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU(inplace=True)
  )
  (layer3): Sequential(
    (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU(inplace=True)
  )
  (fc): Sequential(
    (0): Linear(in_features=4096, out_features=10, bias=True)
  )
)
layer1 Sequential(
  (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
  (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (2): ReLU(inplace=True)
)
layer1.0 Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
layer1.1 BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
layer1.2 ReLU(inplace=True)
layer2 Sequential(
  (0): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
  (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (2): ReLU(inplace=True)
)
layer2.0 Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
layer2.1 BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
layer2.2 ReLU(inplace=True)
layer3 Sequential(
  (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
  (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (2): ReLU(inplace=True)
)
layer3.0 Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
layer3.1 BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
layer3.2 ReLU(inplace=True)
fc Sequential(
  (0): Linear(in_features=4096, out_features=10, bias=True)
)
fc.0 Linear(in_features=4096, out_features=10, bias=True)

각 모듈을 불러왔다. 여기서는 sequential로 묶어서 불렀던 module들도 하나씩 다시 쪼개서 반환해주기 때문에 sequential이 사용된 모델에 접근할 때에는 이 방법이 도움이 될 것이다.

 

 

사용한 모델

class ToyNet(nn.Module):
  def __init__(self):
    super(ToyNet, self).__init__()
    self.layer1 = nn.Sequential(
        nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1, bias=False),
        nn.BatchNorm2d(16),
        nn.ReLU(inplace=True)
    )
    self.layer2 = nn.Sequential(
        nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=2, padding=1, bias=False),
        nn.BatchNorm2d(32),
        nn.ReLU(inplace=True)
    )
    self.layer3 = nn.Sequential(
        nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2, padding=1, bias=False),
        nn.BatchNorm2d(64),
        nn.ReLU(inplace=True)
    )
    self.fc = nn.Sequential(
        nn.Linear(8*8*64, 10)
    )
    for m in self.modules():
      if isinstance(m, nn.Conv2d):
        n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
        m.weight.data.normal_(0, math.sqrt(2./n))
      elif isinstance(m, nn.BatchNorm2d):
        m.weight.data.fill_(1)
        m.bias.data.zero_()

  def forward(self, x):
    out = self.layer1(x) #32x32
    out = self.layer2(out) #16x16
    out = self.layer3(out) #8x8

    out = out.view(x.size(0), -1)
    out = self.fc(out)
    
    return out
    
    net = ToyNet()

 

정리되어 있지 않은 colab notebook 링크

https://colab.research.google.com/drive/1kOrbf-sB-uEWIhQ8Jds55l09kW8_zJZP?usp=sharing