Pytorch로 TAKD 구현해보기

Teacher Assistants Knowledge Distillation

 

선생 모델과 학생모델의 큰 격차로 인해 지식증류가 잘 시행되지 않는 현상을 해결하고자 만든 모델, TAKD입니다.

두 모델 사이에 보조 모델(Teacher Assistant)를 배치하여 지식 증류가 잘 일어나도록 돕습니다.

 

파이썬의 Pytorch 라이브러리를 이용하여 간?단하게 구현할 수 있습니다.

 

사실 논문에 첨부된 깃헙 링크가 있기는 합니다. 살펴보면 NNI라는 툴을 이용하여 구현했습니다.

NNI는 쉽게 말하면 하이퍼파라미터 튜닝을 쉽게 해주는 툴이라고 보시면 됩니다.

 

그런데 제가 이걸 돌려보려고 하니 모델을 아무리 작게 만들어도 도저히 실행이 되질 않더군요..

이런저런 방법을 모두 시도해봤지만 실패해서 결국은 NNI를 사용하지 않고 하이퍼파라미터를 임의로 고정해서 돌려보기로 했습니다.

 

---

 

환경

Google의 Colab이나 GPU가 있는 데스크탑에서 실행 가능합니다. Colab의 경우 무료버전에도 사용 가능한 GPU를 할당해주기 때문에 부담없이 사용할 수 있습니다. 다만 할당량이 금방 바닥나는 편입니다.

 

GPU대신 CPU로도 실행은 가능하지만 실행 시간이 무지막지하게 차이나므로 웬만해선 CPU는 사용하지 않는 것이 좋습니다.

 

GPU를 사용하기 위해선 NVIDIA에서 제공하는 CUDA 시스템을 사용해야 합니다. Colab에는 자동으로 설치가 되어있고, 데스크탑 환경에서 실행하려면 따로 설치가 필요합니다. Pytorch역시 cuda를 지원하는 버전을 설치하셔야 합니다.

CUDA Toolkit 12.1 Downloads

$ pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu121

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

 

---

 

모델 및 데이터셋

훈련 및 테스트에 사용할 데이터셋은 Cifar10 or Cifar100을 사용합니다. torchvision 라이브러리에서 쉽게 얻을 수 있습니다.

def data_loader(num_classes=10):
    # Below we are preprocessing data for CIFAR-10. We use an arbitrary batch size of 128.
    transforms_cifar = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            #데이터셋을 정규화해주기 위한 코드
        ]
    )

    if num_classes == 10:
        # Loading the CIFAR-10 dataset:
        train_dataset = datasets.CIFAR10(
            root="./data", train=True, download=True, transform=transforms_cifar
        )
        test_dataset = datasets.CIFAR10(
            root="./data", train=False, download=True, transform=transforms_cifar
        )

    elif num_classes == 100:
        train_dataset = datasets.CIFAR100(
            root="./data", train=True, download=True, transform=transforms_cifar
        )
        test_dataset = datasets.CIFAR100(
            root="./data", train=False, download=True, transform=transforms_cifar
        )
        
    return train_dataset, test_dataset

 

모델은 논문에선 ResNet과 CNN 두 가지로 소개되었습니다. 비교적 가볍게 돌려보기 위해서 CNN을 사용하도록 하겠습니다.

선생 모델과 학생 모델, 보조 모델의 계층 수를 알맞게 정하는 것이 중요합니다. TAKD는 계층 수가 많이 차이나는 두 모델 사이의 문제를 해결하기 위한 방법이므로 크게 차이가 나지 않으면 의미가 없어집니다. 논문에서 소개한 예시들 중 가장 작은 크기인 10 (T) -> 4 (TA) -> 2 (S)를 사용하도록 하겠습니다.

 

모델을 구현하는 코드는 논문에 소개된 깃헙의 코드 중 모델명만으로 모델을 생성할 수 있게 해주는 ConvNetMaker 클래스가 있어서 이를 가져다가 사용했습니다.

import torch.nn as nn


class ConvNetMaker(nn.Module):
    """
    Creates a simple (plane) convolutional neural network
    """

    def __init__(self, layers):
        """
        Makes a cnn using the provided list of layers specification
        The details of this list is available in the paper
        :param layers: a list of strings, representing layers like ["CB32", "CB32", "FC10"]
        """
        super(ConvNetMaker, self).__init__()
        self.conv_layers = []
        self.fc_layers = []
        h, w, d = 32, 32, 3
        previous_layer_filter_count = 3
        previous_layer_size = h * w * d
        num_fc_layers_remained = len([1 for l in layers if l.startswith("FC")])
        for layer in layers:
            if layer.startswith("Conv"):
                filter_count = int(layer[4:])
                self.conv_layers += [
                    nn.Conv2d(
                        previous_layer_filter_count,
                        filter_count,
                        kernel_size=3,
                        padding=1,
                    ),
                    nn.BatchNorm2d(filter_count),
                    nn.ReLU(inplace=True),
                ]
                previous_layer_filter_count = filter_count
                d = filter_count
                previous_layer_size = h * w * d
            elif layer.startswith("MaxPool"):
                self.conv_layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
                h, w = int(h / 2.0), int(w / 2.0)
                previous_layer_size = h * w * d
            elif layer.startswith("FC"):
                num_fc_layers_remained -= 1
                current_layer_size = int(layer[2:])
                if num_fc_layers_remained == 0:
                    self.fc_layers += [
                        nn.Linear(previous_layer_size, current_layer_size)
                    ]
                else:
                    self.fc_layers += [
                        nn.Linear(previous_layer_size, current_layer_size),
                        nn.ReLU(inplace=True),
                    ]
                previous_layer_size = current_layer_size

        conv_layers = self.conv_layers
        fc_layers = self.fc_layers
        self.conv_layers = nn.Sequential(*conv_layers)
        self.fc_layers = nn.Sequential(*fc_layers)

    def forward(self, x):
        x = self.conv_layers(x)
        x = x.view(x.size(0), -1)
        x = self.fc_layers(x)
        return x

plane_cifar10_book = {
	'2': ['Conv16', 'MaxPool', 'Conv16', 'MaxPool', 'FC10'],
	'4': ['Conv16', 'Conv16', 'MaxPool', 'Conv32', 'Conv32', 'MaxPool', 'FC10'],
	'6': ['Conv16', 'Conv16', 'MaxPool', 'Conv32', 'Conv32', 'MaxPool', 'Conv64', 'Conv64', 'MaxPool', 'FC10'],
	'8': ['Conv16', 'Conv16', 'MaxPool', 'Conv32', 'Conv32', 'MaxPool', 'Conv64', 'Conv64', 'MaxPool',
		  'Conv128', 'Conv128','MaxPool', 'FC64', 'FC10'],
	'10': ['Conv32', 'Conv32', 'MaxPool', 'Conv64', 'Conv64', 'MaxPool', 'Conv128', 'Conv128', 'MaxPool',
		   'Conv256', 'Conv256', 'Conv256', 'Conv256' , 'MaxPool', 'FC128' ,'FC10'],
}


plane_cifar100_book = {
	'2': ['Conv32', 'MaxPool', 'Conv32', 'MaxPool', 'FC100'],
	'4': ['Conv32', 'Conv32', 'MaxPool', 'Conv64', 'Conv64', 'MaxPool', 'FC100'],
	'6': ['Conv32', 'Conv32', 'MaxPool', 'Conv64', 'Conv64', 'MaxPool','Conv128', 'Conv128' ,'FC100'],
	'8': ['Conv32', 'Conv32', 'MaxPool', 'Conv64', 'Conv64', 'MaxPool', 'Conv128', 'Conv128', 'MaxPool',
		  'Conv256', 'Conv256','MaxPool', 'FC64', 'FC100'],
	'10': ['Conv32', 'Conv32', 'MaxPool', 'Conv64', 'Conv64', 'MaxPool', 'Conv128', 'Conv128', 'MaxPool',
		   'Conv256', 'Conv256', 'Conv256', 'Conv256' , 'MaxPool', 'FC512', 'FC100'],
}

 

쉽게 설명하면, 원하는 데이터셋에 해당하는 딕셔너리(현재는 plane_cifar100_book)에서 계층의 갯수를 가지고 해당 계층에 딱 맞는 hidden layer들을 이름 그대로 뽑아옵니다.

뽑아온 이름들을 ConvNetMaker에 넣기만 하면 알아서 모델이 만들어지는 엄청난 갓코드입니다.

layer = plane_cifar100_book.get("10")
teacher_model = ConvNetMaker(layer).to(device)

layer = plane_cifar100_book.get("4")
TA_model = ConvNetMaker(layer).to(device)

layer = plane_cifar100_book.get("2")
takd_student_model = ConvNetMaker(layer).to(device)

 

구현

기본이 되는 세팅들을 모두 구현했으면 이제 훈련과 테스트를 진행할 차례입니다.

주요한 하이퍼파라미터들은 다음과 같이 설정하였습니다.

• Learning rate: 0.1
• batch size: 128
• epoch: 160
• SGD 사용
• momentum: 0.9, nesterov
• weight decay: 1e-4

Temperature, Soft_target_loss_weight, ce_loss_weight의 경우는 본래 NNI를 이용하여 튜닝해야 하지만 여기서는 NNI를 사용하지 않기 때문에 각각 2, 0.25, 0.75로 임의 고정해서 사용하도록 하겠습니다.

 

학습 순서는 다음과 같습니다.

선생 모델의 학습 -> 선생 모델로부터 보조 모델 지식증류 -> 보조 모델로부터 학생 모델 지식증류

train(teacher_model, train_loader, epochs=160, learning_rate=0.1, device=device)
test_accuracy_deep = test(teacher_model, test_loader, device)

train_knowledge_distillation(
        teacher=teacher_model,
        student=TA_model,
        train_loader=train_loader,
        epochs=160,
        learning_rate=0.1,
        T=2,
        soft_target_loss_weight=0.25,
        ce_loss_weight=0.75,
        device=device,
    )
train_knowledge_distillation(
        teacher=TA_model,
        student=takd_student_model,
        train_loader=train_loader,
        epochs=160,
        learning_rate=0.1,
        T=2,
        soft_target_loss_weight=0.25,
        ce_loss_weight=0.75,
        device=device,
    )

test_accuracy_light_ce_and_takd = test(takd_student_model, test_loader, device)

 

결과

Teacher accuracy: 51.82%
Student accuracy without teacher: 26.77%	#학생 모델의 단순학습 결과
Student accuracy with CE + KD: 31.34%		#단순 지식증류의 결과
Student accuracy with CE + TAKD: 36.71%		#TAKD의 결과

 

모델의 크기 자체가 작아 정확도가 그렇게 높진 않지만 확연한 차이는 보입니다. 다만, 일반 학습의 정확도가 논문에 소개된 것에 비해 현저히 작은 수치를 보이는데, 왜 그런지는 더 탐구해봐야 할 것 같습니다. 아무래도 NNI를 이용한 하이퍼파라미터 튜닝의 여부가 중요한 듯 싶습니다.

 

 

 

GitHub - YSW2/CV-KnowledgeDistillation