ResNet처럼 Batch Normalization을 포함한 네트워크를 Feature Extractor로 사용할 때 유의할 점

  ResNet과 같이 배치 정규화(batch normalization)를 포함하고 있는 네트워크를 특징 추출기(feature extractor)로 사용할 때 유의해야 할 점이 있다. 별로 중요하지 않은 것처럼 보여도, 실제로 모델을 만드는 입장에서 제대로 이해하고 있지 않으면 많이 헤맬 수 있는 부분이다.

 

  1. 학습 모드(training mode)와 평가 모드(evaluation mode)일 때 추출되는 특징 맵(feature map)이 다르다는 점

 

  기본적으로 배치 정규화는 학습 시 사용하는 파라미터와 평가 시 사용하는 파라미터가 다르게 적용된다. 따라서 만약 batch normalization을 포함한 네트워크를 feature extractor로 사용하고자 한다면, 같은 이미지에 대해 매번 동일한 특징 맵이 추출될 수 있도록 하기 위해 항상 train() 모드로 사용하거나 항상 eval() 모드로 사용하는 식으로 일관적일 필요가 있다. 일반적인 목적의 특징 추출기로 사용한다면 eval() 모드로만 사용하는 것을 추천한다. (학습할 때 train(), 평가할 때 eval()을 사용하는 것도 일반적인 분류 모델 학습 목적이라면 성능에 문제가 없다.)

 

  전이 학습(transfer learning)을 수행할 때 ResNet 기반의 고정된 특징 추출기가 필요하다면, conv 레이어 부분을 eval() 모드로 사용하면 된다. 만약 이를 지키지 않고 일관적이지 않게 train()과 eval()을 번갈아 호출하게 된다면, 동일한 이미지에 대하여 추출되는 feature maps이 매번 다른 값을 가질 수 있다. 앞서 언급했듯이 일반적인 분류 문제에서는 이게 큰 이슈가 되지 않는다. 다만 동일 이미지에 대한 feature maps가 항상 같아야 되는 경우에는 문제가 된다.

 

  2. 다른 네트워크를 포함한 네트워크에서 train()이나 eval()을 사용하는 경우

 

  필요한 경우 특징 추출기(feature extractor)를 포함한 하나의 네트워크 자체를 새롭게 정의할 수 있다. 이는 학습 코드 구현상의 편리성을 주는 경우가 많기 때문에, 종종 볼 수 있는 코드 유형이다. 예를 들어 ResNet18 모델의 앞부분은 고정한(fixed) 상태로 뒤쪽 FC 레이어만 새롭게 교체하여 학습하는 전이 학습(transfer learning) 방법을 사용할 수 있다. 소스코드 예시는 다음과 같다.

 

class StudentNetwork(nn.Module):
    def __init__(self):
        super(StudentNetwork, self).__init__()

        self.feature_extractor = nn.Sequential(*list(models.resnet18(pretrained=True).children())[:-1]).eval()
        self.fc = nn.Linear(512, 2) # binary classification (num_of_class == 2)

        # fix the pre-trained network
        for param in self.feature_extractor.parameters():
            param.requires_grad = False

    def forward(self, images):
        features = self.feature_extractor(images)
        x = torch.flatten(features, 1)
        outputs = self.fc(x)
        return features, outputs

 

  만약 위와 같이 모델을 정의했다면, 특징 추출기(feature extractor)는 항상 eval() 모드로 수행되는 것이다. 따라서 StudentNetwork 클래스의 인스턴스에 대하여 별도로 train()이나 eval()을 호출할 필요가 없다. 그냥 학습 단계든 평가 단계든 그 단계(phase)에 상관없이 그대로 사용하면 된다. 그러면 동일한 이미지에 대하여 항상 같은 features가 반환된다.

 

  위와 같이 ResNet 기반의 특징 추출기를 내부적으로 포함하고 있는 하나의 모델 인스턴스 model을 초기화한 경우를 생각해보자. 만약에 이때 model.train()을 호출하면 내부에 포함된 ResNet에 대해서도 학습 모드(train mode)가 적용되기 때문에, 의도치 않게 특징 추출기에 영향을 미칠 수 있다. 이 경우 동일한 이미지에 대하여 forward()의 결과인 features 텐서 값이 변경될 수 있는 것이다. 이 부분은 특히나 실수하기 쉬운 부분이므로 유의하자. 

 

  [참고] PyTorch에서 모델을 초기화하면 기본 설정으로 requires_grad 값이 True가 된다. (이는 사전 학습된 네트워크를 불러올 때에도 마찬가지다!) 따라서 별도로 명시하지 않는다면 자동으로 기울기(gradient)를 추적하기 때문에, 학습하지 않고자 하는 레이어에 대해서는 requires_grad 값을 명시적으로 False로 설정할 필요가 있다.