PyTorch 모델 특정 종류의 레이어를 바꾸기

API를 쓰다보면, activation function을 비롯해서 모델의 레이어를 바꾸고 싶은 일이 종종 생기는데, nn.Sequantial()이 다중으로 겹쳐진 경우 구글링 해서 나오는 코드들은 대부분 동작을 안 한다. (예시 isinstance()를 한두번 정도 묻는 것, 혹은 setattr로 변경을 시도하는 것)

코드

아래의 코드는 재귀함수 형태로, 바꾸고 싶은 instance를 target에 전달하고, 바꾸고 싶은 instance를 source에 전달을 한다. 이러면 nn.Sequantial에 얼마나 감싸져있던 재귀형태로 모두 진행하여 모델 레이어를 변경할 수 있다.

def replace_module(modules:nn.Module, target, source):
    for name, child in modules.named_children():
        if isinstance(child, target):
            modules._modules[name] = source
        elif isinstance(child, nn.Sequential):
            replace_module(child, target, source)

사용예시

unet 모델의 LeakyReLU를 전부 SiLU (SWISH)로 변경하는 예제

target = nn.LeakyReLU
source = nn.SiLU()
model = unet 

def replace_module(modules:nn.Module, target, source):
    for name, child in modules.named_children():
        if isinstance(child, target):
            modules._modules[name] = source
        elif isinstance(child, nn.Sequential):
            replace_module(child, target, source)

replace_module(model, target, source)

결과

BasicUNet(
  (conv_0): TwoConv(
    (conv_0): Convolution(
      (conv): Conv3d(1, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
      (adn): ADN(
        (N): InstanceNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
        (D): Dropout(p=0.0, inplace=False)
        (A): LeakyReLU(negative_slope=0.1, inplace=True)
      )
    )
    (conv_1): Convolution(
      (conv): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
      (adn): ADN(
        (N): InstanceNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
        (D): Dropout(p=0.0, inplace=False)
        (A): LeakyReLU(negative_slope=0.1, inplace=True)
      )
    )
  )
  (down_1): Down(
    (max_pooling): MaxPool3d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (convs): TwoConv(
      (conv_0): Convolution(
        (conv): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
      (conv_1): Convolution(
        (conv): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
    )
  )
  (down_2): Down(
    (max_pooling): MaxPool3d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (convs): TwoConv(
      (conv_0): Convolution(
        (conv): Conv3d(32, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
      (conv_1): Convolution(
        (conv): Conv3d(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
    )
  )
  (down_3): Down(
    (max_pooling): MaxPool3d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (convs): TwoConv(
      (conv_0): Convolution(
        (conv): Conv3d(64, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
      (conv_1): Convolution(
        (conv): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
    )
  )
  (down_4): Down(
    (max_pooling): MaxPool3d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (convs): TwoConv(
      (conv_0): Convolution(
        (conv): Conv3d(128, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
      (conv_1): Convolution(
        (conv): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
    )
  )
  (upcat_4): UpCat(
    (upsample): UpSample(
      (deconv): ConvTranspose3d(256, 128, kernel_size=(2, 2, 2), stride=(2, 2, 2))
    )
    (convs): TwoConv(
      (conv_0): Convolution(
        (conv): Conv3d(256, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
      (conv_1): Convolution(
        (conv): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
    )
  )
  (upcat_3): UpCat(
    (upsample): UpSample(
      (deconv): ConvTranspose3d(128, 64, kernel_size=(2, 2, 2), stride=(2, 2, 2))
    )
    (convs): TwoConv(
      (conv_0): Convolution(
        (conv): Conv3d(128, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
      (conv_1): Convolution(
        (conv): Conv3d(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
    )
  )
  (upcat_2): UpCat(
    (upsample): UpSample(
      (deconv): ConvTranspose3d(64, 32, kernel_size=(2, 2, 2), stride=(2, 2, 2))
    )
    (convs): TwoConv(
      (conv_0): Convolution(
        (conv): Conv3d(64, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
      (conv_1): Convolution(
        (conv): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
    )
  )
  (upcat_1): UpCat(
    (upsample): UpSample(
      (deconv): ConvTranspose3d(32, 32, kernel_size=(2, 2, 2), stride=(2, 2, 2))
    )
    (convs): TwoConv(
      (conv_0): Convolution(
        (conv): Conv3d(64, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
      (conv_1): Convolution(
        (conv): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
    )
  )
  (final_conv): Conv3d(32, 2, kernel_size=(1, 1, 1), stride=(1, 1, 1))
)

이후

BasicUNet(
  (conv_0): TwoConv(
    (conv_0): Convolution(
      (conv): Conv3d(1, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
      (adn): ADN(
        (N): InstanceNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
        (D): Dropout(p=0.0, inplace=False)
        (A): SiLU()
      )
    )
    (conv_1): Convolution(
      (conv): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
      (adn): ADN(
        (N): InstanceNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
        (D): Dropout(p=0.0, inplace=False)
        (A): SiLU()
      )
    )
  )
  (down_1): Down(
    (max_pooling): MaxPool3d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (convs): TwoConv(
      (conv_0): Convolution(
        (conv): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): SiLU()
        )
      )
      (conv_1): Convolution(
        (conv): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): SiLU()
        )
      )
    )
  )
  (down_2): Down(
    (max_pooling): MaxPool3d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (convs): TwoConv(
      (conv_0): Convolution(
        (conv): Conv3d(32, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): SiLU()
        )
      )
      (conv_1): Convolution(
        (conv): Conv3d(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): SiLU()
        )
      )
    )
  )
  (down_3): Down(
    (max_pooling): MaxPool3d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (convs): TwoConv(
      (conv_0): Convolution(
        (conv): Conv3d(64, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): SiLU()
        )
      )
      (conv_1): Convolution(
        (conv): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): SiLU()
        )
      )
    )
  )
  (down_4): Down(
    (max_pooling): MaxPool3d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (convs): TwoConv(
      (conv_0): Convolution(
        (conv): Conv3d(128, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): SiLU()
        )
      )
      (conv_1): Convolution(
        (conv): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): SiLU()
        )
      )
    )
  )
  (upcat_4): UpCat(
    (upsample): UpSample(
      (deconv): ConvTranspose3d(256, 128, kernel_size=(2, 2, 2), stride=(2, 2, 2))
    )
    (convs): TwoConv(
      (conv_0): Convolution(
        (conv): Conv3d(256, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
      (conv_1): Convolution(
        (conv): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
    )
  )
  (upcat_3): UpCat(
    (upsample): UpSample(
      (deconv): ConvTranspose3d(128, 64, kernel_size=(2, 2, 2), stride=(2, 2, 2))
    )
    (convs): TwoConv(
      (conv_0): Convolution(
        (conv): Conv3d(128, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
      (conv_1): Convolution(
        (conv): Conv3d(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
    )
  )
  (upcat_2): UpCat(
    (upsample): UpSample(
      (deconv): ConvTranspose3d(64, 32, kernel_size=(2, 2, 2), stride=(2, 2, 2))
    )
    (convs): TwoConv(
      (conv_0): Convolution(
        (conv): Conv3d(64, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
      (conv_1): Convolution(
        (conv): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
    )
  )
  (upcat_1): UpCat(
    (upsample): UpSample(
      (deconv): ConvTranspose3d(32, 32, kernel_size=(2, 2, 2), stride=(2, 2, 2))
    )
    (convs): TwoConv(
      (conv_0): Convolution(
        (conv): Conv3d(64, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
      (conv_1): Convolution(
        (conv): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1))
        (adn): ADN(
          (N): InstanceNorm3d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)
          (D): Dropout(p=0.0, inplace=False)
          (A): LeakyReLU(negative_slope=0.1, inplace=True)
        )
      )
    )
  )
  (final_conv): Conv3d(32, 2, kernel_size=(1, 1, 1), stride=(1, 1, 1))
)

'개발새발 > 개발 셋업' 카테고리의 다른 글

selenium을 활용한 원본 이미지 크롤링 파이썬 코드 (1)	2023.04.18
Segmentation 수행을 위한 PyTorch-lightning 설정 예제 (0)	2023.03.16
윈도우 포맷 후 환경 설정 및 설치 체크리스트 (0)	2023.02.24
PyTorch Lightning 소개 및 설명 (0)	2023.02.13
[Docker, Linux] GPG에러 해결, public key 수동 설치 (apt-get이 에러가 발생하며 동작하지 않을 때) (0)	2023.01.19

사용예시

'개발새발 > 개발 셋업' 카테고리의 다른 글

검색 태그

티스토리툴바