[졸업프로젝트] Similarity based knowledge distillation method for light camera pose regressor
지난 [졸업프로젝트] PoseNet 경량화를 통한 Jetson Nano RC카의 실내 자율주행 구현 이후의 프로젝트 진행 과정에 대한 포스팅입니다.
저번 2022-2에 저희 팀은 Camera Pose Estimation 모델인 PoseNet을 Knowledge Distillation 을 통해 경량화하는 프로젝트를 진행했습니다. 프로젝트를 진행하며 다음의 프로젝트의 변경사항이 있었습니다.
- GoogLeNet Backbone인 PoseNet이 아닌 ResNet을 Backbone으로 하는 Camera Pose Estimation 모델을 만들어 해당 모델을 경량화하는 프로젝트로 변경했습니다. 이 이유는 Knowledge Distillation 의 경우 Distillation 이 Teacher와 Student 모델의 module 단위로 이루어져야 하는데, GoogLeNet의 경우 이러한 module이 구분되지 않기 때문입니다. 실제 저희가 실험을 위한 ResNet기반 pose regressor 모델을 만들었고, 해당 모델이 기존의 GoogLeNet기반의 PoseNet보다 성능이 좋은 것을 확인했습니다.
PoseRegressor 경량화 과정
Knowledge Distillation 이란?
우선 저희가 사용하는 경량화 기법에 대해 설명하겠습니다.
[출처: 2021: Knowledge Distillation Survey]
KD(Knowledge Distillation) 이란 성능이 좋고, 모델 사이즈가 큰 Teacher 모델의 지식을 상대적으로 성능이 나쁘고, 모델 사이즈가 작은 Student 모델에 전달하는 경량화 기법입니다. 위 그림은 Vinilla KD인 Response-Based Knowledge Distillation 에 대한 그림입니다.
KD에서 중요한 부분은 1) Teacher의 어떤 지식을 2) 어떻게 Student에게 전달할지 입니다. 이 1), 2) 의 방법에 따라 다양한 종류의 KD 기법이 있습니다.
[출처: 2021: Knowledge Distillation Survey]
Regression Problem
처음 제시된 Vinilla KD와 위의 방법론들 대부분 DNN의 Classification Task에 초점을 맞춰 발전되어왔습니다. Classification 의 경우 prediction 값을 softmax를 통해 확률값으로 전달할 수 있습니다.
그러나 저희 프로젝트의 Camera Pose Regression 은 연속적인 값을 다루는 Regression task 입니다. Regression problem은 prediction 값에 softmax가 적용되지 않고, class에 대한 확률이 없어 Teacher 의 prediction 값을 전달하는 방식은 적절하지 않습니다. 따라서 저희는 camera pose estimation 에 적합한 저희만의 방법론을 생각해야 했습니다.
저희는 위 KD 기법 중 Feature-Based KD 와 Relation-Based KD를 사용했습니다. Camera Pose estimation 은 공간에 관한 이미지와 6-DOF 값을 다룹니다. 각 Dataset의 이미지에 대한 관계, 각 Feature map의 관계, ground truth의 6-DOF관계를 전달하는 방식으로 Teacher의 Knowledge를 전달했습니다.
각 요소들의 관계는 Pairwise Cosine Similarity 를 이용했습니다.
def Pairwise_Cosine_Similarity(input_m):
# shape : N*256*256*3 => N*(256*256*3)
sim = []
for feature in input_m:
temp = (feature.view([feature.shape[0], -1]).squeeze())
# F.pdist => 1*(_nC_2) //
sim.append(pairwise_cosine_similarity(temp)) # N*N
return sim또한 실제로 pairwise similarity 값을 plot해보며 저희 방법에 대한 시각적인 검증을 진행했습니다. (아래 그림은 head dataset에 대한 image간 pairwise similarity, Ground truth 값 간 pairwise similarity입니다.)
경량화 과정
Model
저희는 ResNet 을 기반으로 하는 Camera Pose estimation model들을 pytorch 로 구현했습니다. Classification 이 아닌 position 과 orientation에 관한 7dimension(pos(x,y,z) + quaternion(w,p,q,r)) 을 얻기 위해 마지막 fc layer가 7x1 값을 출력하도록 수정하였습니다. 아래 코드는 모델 수정 과정의 예시입니다.
# defines the regression module for resnet
class Regression(nn.Module):
def __init__(self, weights=None, input_cls=512, mid_cls=1024):
super(Regression, self).__init__()
dropout_rate = 0.5
self.dropout = nn.Dropout(p=dropout_rate)
#self.projection = nn.AvgPool2d(kernel_size=7, stride=1)
self.cls_fc_pose = nn.Sequential(*[weight_init_resnet("pose", nn.Linear(input_cls, mid_cls)),
nn.ReLU(inplace=True)])
self.cls_fc_xy = weight_init_resnet("XYZ", nn.Linear(mid_cls, 3), weights=weights)
self.cls_fc_wpqr = weight_init_resnet("WPQR", nn.Linear(mid_cls, 4), weights=weights)
def forward(self, input):
output = self.cls_fc_pose(input.view(input.size(0), -1))
output = self.dropout(output)
output_xy = self.cls_fc_xy(output)
output_wpqr = self.cls_fc_wpqr(output)
output_wpqr = F.normalize(output_wpqr, p=2, dim=1)
return [output_xy, output_wpqr]def resnetxxim(pretrained=False, progress=True, isKD=False, isTest=False):
# Pretrained = True : ImageNet
resnetxx = _resnet('resnet34im', BasicBlock, [_, _, _, _], pretrained, progress, isKD)
resnetxx.fc = Regression()
return resnetxx이렇게 수정된 resnet 기반 Camera Pose Regressor를 이제부터 resPoseNet이라 말하겠습니다.
저희는 Teacher model로 ResNet34, ResNet50 을 선택하였고, Student model 로 ResNet18을 선택하였습니다.
각 모델의 Parameter수는 다음과 같습니다.
resnet18 : 46 MB
resnet34: 87.27MB
resnet50: 102.45MB
저희 KD Architecture 를 적용할 경우 resnet18이 KD를 사용하기 이전보다 성능이 좋아지고, Teacher 모델의 성능을 능가하기도 한다는 것을 확인했습니다.
실험환경
KD 실험은 google colab 환경에서 진행하였습니다. 모델 성능 monitoring은 wandb(weight&bias) 를 사용했습니다.
Dataset
Camera Pose Estimation의 Dataset은 실외, 실내 데이터로 나눌 수 있습니다. 이중 가장 대표적인 Benchmark dataset을 사용했습니다. 각 dataset의 특징은 다음과 같습니다.
실내 : 7Scenes Dataset 중 Heads, Fire, Stairs
실외 : Kings College, Old Hospital, Shop Facade
모델학습 및 실험
모델학습 및 실험에 사용한 코드들 및 hierarchy 입니다.
train.py test.py 는 KD를 적용하지 않은 파이썬 파일입니다.
KD_train.py KD_test.py 는 KD 학습 파이썬 파일입니다.
Teacher Model 학습
!python train.py --model [resnet18 | resnet34 | resnet50 | resnet101] --dataroot [DATAROOT] --name [모델명]/[Dataset]/[beta500_ex] --beta 500 --gpu 0 --niter 500 --batchSize32 --lr 0.001colob 또는 terminal에서 위 command로 학습을 시작합니다. --model 옵션에서 학습시킬 모델의 종류, --dataroot에서 dataset의 root, --name으로 학습된 .pth 파일이 저장될 디렉토리 명을 지정할 수 있습니다. --gpu 를 통해 사용할 gpu id 를 넣어줍니다.
!python test.py --model [resnet18 | resnet34 | resnet50 | resnet101] --dataroot [DATAROOT] --name [모델명]/[Dataset]/[beta500_ex] --gpu 0
위 코드를 통해 학습된 모델을 test/ validation 합니다.
Student Model 학습
!python KD_train.py --model [resnet18 | resnet34 | resnet50 | resnet101] --dataroot [DATAROOT] --name [S_모델명]/[Dataset]/[beta500_bt_lr_m#1_m#2_scaling] --beta 500 --gpu 0 --niter 500 --T_model [ resnet34 | resnet50 | resnet101] --T_path [TeacherModel_Path] --save_epoch_freq 5 --CSmodule 3 4 --hintmodule 5 [--CSKL]KD_train.py를 위 명령어로 수행합니다. Teacher 모델 및 여러 저희 architecture의 KD 옵션을 사용할 수 있습니다.
!python KD_test.py --model [resnet18 | resnet34 | resnet50 | resnet101] --dataroot [DATAROOT] --name Student/[Dataset]/[beta500_LossFunction] --beta 500 --gpu 0
KD_test.py를 위 명령어로 수행합니다.
학습 과정 및 test 결과에 대한 자세한 과정을 아래 그림과 같은 wandb project에서 모니터링하며 tuning 을 진행했습니다.
Jetson Nano/ 저전력 환경 모델
앞선 과정으로 만들어진 저희는 경량화된 모델들을 얻었습니다. 이제 해당 모델들이 5W전력의 Jetson nano에서 잘 수행되는지 확인했습니다.
Jetson Nano Version 및 로컬 환경 구성
우선 저희가 사용한 Jetson Nano 와 JetPack 버전입니다.
torch를 설치하기 위해선 JetPack 버전과 Cuda, Cudnn 버전을 꼭 확인하여 일치시켜야 합니다.
현재 위 사진에서는 CUDA : NO 라고 되어있는대요, 실제로 CUDA 가 설치되어 있지만 path 문제로 인식되지 않는 상황입니다.
아래와 같이 코드를 추가하면 CUDA가 인식되고, gpu를 사용할 수 있습니다.
이제 requirements.txt 를 통해 필요한 패키지를 설치하고, cuda에 맞는 torch를 설치합니다. jetson nano의 경우 torch를 아래 링크에서 버전에 맞는 파일을 다운받아 설치해야 합니다.https://forums.developer.nvidia.com/t/pytorch-for-jetson/72048
Jetson nano에서 모델 돌리기
환경 구성이 끝나고, terminal에서 모델이 잘 돌아가는 것을 확인했습니다.
저전력 환경에서의 성능 (jtop)
jetson 시리즈에는 jtop 이라고 시스템을 모니터링 할 수 있는 프로그램이 있습니다. jtop을 설치 후 jtop 을 터미널에 명령할 경우 다음과 같은 모니터링 화면을 볼 수 있습니다.
또한 jtop api를 통해 각 cpu, gpu 사용량, cpu,gpu의 temperature 등을 저장할 수 있습니다.
jtop_logger.py jtop 깃허브에 있는 api sample를 사용했습니다.
from jtop import jtop, JtopException
import csv
import argparse
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Simple jtop logger')
# Standard file to store the logs
parser.add_argument('--file', action="store", dest="file", default="log.csv")
args = parser.parse_args()
print("Simple jtop logger")
print("Saving log on {file}".format(file=args.file))
try:
with jtop() as jetson:
# Make csv file and setup csv
with open(args.file, 'w') as csvfile:
stats = jetson.stats
# Initialize cws writer
writer = csv.DictWriter(csvfile, fieldnames=stats.keys())
# Write header
writer.writeheader()
# Write first row
writer.writerow(stats)
# Start loop
while jetson.ok():
stats = jetson.stats
# Write row
writer.writerow(stats)
print("Log at {time}".format(time=stats['time']))
except JtopException as e:
print(e)
except KeyboardInterrupt:
print("Closed with CTRL-C")
except IOError:
print("I/O error")
# EOF
다음과 같은 CSV 파일을 얻을 수 있다.
현재 jetson nano 에서 모델을 돌려 해당 csv 파일 데이터를 수집하고, plot 하는 과정에 있습니다.
이외 각 모델별 5W 전력 Jetson nano에서 inference time의 경우 다음과 같습니다.
-inference time
Resnet18 : 50
Resnet34 : 90
Resnet50 : 125
inference time은 다음 코드를 이용해 측정했습니다.
결론
아직까지 남은 실험이 있으나 현재까지 대부분의 데이터셋에서 저희 모델이 KD를 적용하기 전보다 좋은 결과가 나온 것을 확인했습니다. 아직까지 대부분의 KD는 classification 에서 연구되는데 regression task에서 KD를 적용하는 방법을 제시했다는 것과, 자율주행 등의 분야에서 각광받은 Camera Pose Estimation 모델을 효과적으로 경량화하는 방법을 제시했다는 것에서 이 프로젝트의 의의가 있습니다. 또한 실제 저전력 device인 Jetson nano 에서 경량화된 모델이 더 빠른 inference time으로 좋은 성능을 보이는 것을 확인했습니다.
논문에 자세히 들어갈 내용을 포함시킬 수 없어 자세한 아키텍처와 방법론에 대한 설명은 이 포스팅에 포함되지 않았지만, 이번 프로젝트를 진행하는 과정에 대한 큰 틀을 적었습니다. 
