목표 : 내 데이터셋에 대해 DETR를 재학습시키기
이 글은 DETR에 대한 환경설정이 모두 완료되었다고 가정하고 작성되었습니다. 아직 환경 설정이 완료되지 않았다면 아래 글을 참고해주세요.
Anaconda 설치 및 가상환경 생성
DETR 맛보기
우선 임의의 폴더 상에 DETR 코드를 다운로드합니다. Git bash를 열고 아래 명령을 치면 다운로드가 완료됩니다.
git clone https://github.com/facebookresearch/detr.git .저는 anaconda를 설치할 때 함께 설치되는 spyder(IDE)를 활용했습니다.
- 가상환경을 새로 만들었다면, 환경이 적용되도록 아래 명령을 통해 앱을 새로이 설치해야 합니다.
conda install spyder프롬프트에 아래 명령을 수행하면 가상환경이 적용된 spyder 앱이 실행됩니다.
spyder수정된 스크립트는 아래 github에 보실 수 있습니다.
main.py
- 먼저 main.py를 수정해야합니다.
- 아래는 변경 및 추가된 코드들에 대한 설명입니다.
- torch 라이브러리의 메서드 nn을 불러와야합니다. (나중에 사용)
import torch.nn as nn- argparse 라이브러리는 프로그램을 실행시킬 때 사용자가 원하는 매개변수를 설정하는 데 도움을 줍니다. 아래 코드는 내가 사용하는 데이터셋의 클래스의 개수를 DETR에게 알려주는 역할입니다.
parser.add_argument('--want_class', type=int,
help='number of class which want to finetuning')- 저희는 모델의 모든 parameters를 학습에 사용하지 않습니다. 따라서 사용하지 않는 parameters가 있다는 것을 모델에게 알려주기 위해 아래와 같은 변경이 필요합니다.
# 변경 전
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
# 변경 후
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)- PyTorch는 모델을 로드할 때 저장된 모델 파일에 저장된 모델 구조와 동일한 구조를 가지고 있는 지 확인합니다.
- 하지만 우리는 다른 구조의 모델을 사용할 것이기에 아래와 같이 변경해(strict 옵션을 False로 변경) 이를 알려야합니다.
# 변경 전
model_without_ddp.load_state_dict(checkpoint['model'])
# 변경 후
model_without_ddp.load_state_dict(checkpoint['model'], strict=False)- 가장 끝단의 선형 함수(혹은 fc layer)를 변경합니다.
- 기존에는 coco의 class 개수인 80개(혹은 누락된 값을 포함한 91개)에 대한 분류를 진행하기 때문에 이를 데이터셋에 맞게 변경해야합니다.
- 기존의 layer를 사용하는 데이터셋에 맞는 layer로 바꿔끼었다 정도로 보시면 될 것 같네요.
num_ftrs = model_without_ddp.class_embed.in_features
model_without_ddp.class_embed = nn.Linear(num_ftrs, args.want_class + 2).to(device)models/detr.py
- 아래 코드를 추가해줍니다.
num_classes = args.want_class + 1Preparing Dataset
데이터셋은 COCO format으로 구성되어 있다고 가정합니다.
- 저는 준비된 데이터셋의 폴더 및 파일명을 아래와 같이 변경했습니다.
- train2017 폴더 아래에는 학습에 사용할 이미지들이, val2017 폴더 아래에는 test 및 검증을 위해 사용할 이미지들이 있어야합니다.
- annotations 폴더 아래에는 아래와 같이 학습과 검증용 이미지들에 대한 정보들이 저장된 json 파일이 존재해야합니다.
Finetuning
이제 아래 명령을 통해 새로운 데이터셋을 학습시킬 수 있습니다.
- 학습된 모델은 output_dir 옵션에 지정한 위치에 저장됩니다.
# 단일 GPU를 사용할 경우
python main.py --coco_path /path/to/dataset --resume https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth --want_class number_of_classes_of_dataset --output_dir finetuning_detr
# 다수의 GPU를 사용할 경우
python -m torch.distributed.launch --nproc_per_node=number_of_gpu --use_env main.py --coco_path /path/to/dataset --resume https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth --want_class number_of_classes_of_dataset --output_dir finetuning_detr
AI Explorer