전이 학습으로 골절 이미지 분류하기(1)_ImageFolder

https://www.kaggle.com/datasets/pkdarabi/bone-break-classification-image-dataset

Bone Break Classification Image Dataset

 

https://github.com/ackrilll/TransferLearning

GitHub - ackrilll/TransferLearning

 

몇 일 전 kaggle 에서 재미있는 데이터셋을 발견하였다. 골절 종류를 분류하는 모델을 만들 수 있을 것 같아서 데이터를 다운받아 보았다. 기존에 만들었던 분류 프로젝트와 유사하여 쉬울 것이라 생각했지만 커스텀 데이터셋을 사용하면서 쉽게 해결하지 못했던 문제가 있어 자료로 남겨놓는다.

 

1. ImageFolder 클래스  

처음에는 데이터 셋 구조가 왼쪽 이미지 처럼 되어있었는데 torchvision.datasets 의 ImageFolder 클래스를 사용하기 위해 오른쪽 구조처럼 수정해주었다. 처음에는 디렉토리 구조를 바꾸지 않고 ImageFolder 클래스를 인스턴스화 하려고 하니 반복문을 써서 

데이터를 합쳐줘야 하는 번거로움이 있었다. Roboflow 에서 데이터 셋을 만들면 왼쪽 구조와 같은 데이터 셋이 만들어준다는 것 같다. 

왼쪽은 초기 데이터 셋 디렉토리 구조, 오른쪽은 수정한 데이터 셋 디렉토리 구조

2. 파일 이름 수정

모델을 훈련 시키는 과정에서 파일 이름을 찾을 수 없다는 에러가 나와 확인을 해 본 결과 파일 이름에 확장자 명이 끼어져 있거나 하면 읽을 때 문제가 될 수 있다고 한다. 그래서 각 클래스의 Train, Test 폴더 내부에 존재하는 모든 파일의 이름을 인덱스로 수정해 주었더니 에러를 해결 할 수 있었다. 

 

3. 전이 학습

전이 학습은 하나의 특정 문제에 대해 학습된 모델의 지식(특징 추출 능력, 가중치 등)을 전혀 다른, 하지만 관련성이 있는 다른 문제 해결에 재활용하는 머신러닝 기법이다. 직접 모델을 만들어 사용해도 되지만, 처음부터 모델을 학습시키기 위해서는 너무나 많은 데이터가 필요하기 때문에 대용량 이미지 데이터 ImageNet을 학습시킨 ResNet18 모델을 로드하여 파인 튜닝하였다.

 ImageNet은 1000개의 class를 분류할 때 사용되는 데이터셋이다. 따라서 ResNet18 모델의 fc 레이어는 1000개의 출력 노드를 가지고 있다. 하지만 우리는 10가지 골절 종류에 대한 분류 작업을 수행할 것이므로 fc 레이어를 1000개가 아닌 10개의 출력 노드를 가진 레이어로 바꿔 주었다. 

 

이번 포스팅에서는 모델을 파인 튜닝하여 모델을 추가 학습 시킬 때 발생한 에러에 대해 알아보았다. 사실 너무 디렉토리 구조라는 것이 너무 간단해 보여서 정리를 해야할까 고민했지만, 수정하는데 적지 않은 시간이 걸렸던 것을 생각하면 정리해두는 것이 좋다고 생각이 들었다.