BLOG
인공지능

Transfer Learning을 이용한 Youtube Thumbnail 이미지 카테고리 분류


March 17, 2023, 2:21 p.m.



이번 포스팅에서는 학교 AI 학회 활동의 일환으로 토이 프로젝트를 진행했던 과정에 대해서 포스팅 해보려고 한다.

.

먼저 주제는 기존 베이스라인 모델을 이용해서 Youtube Thumbnail 이미지의 카테고리를 분류하는 것이다.

.

데이터는 팀원이 Youtube Api를 사용해서 이미지 - 카테고리 쌍으로 수집하였다.

.

총 수집한 데이터와 카테고리 별 데이터 수를 plot한 모습이다. 특정 카테고리에 몰려있는 모습을 볼 수 있다. 이 상태로 그대로 학습을 진행하게 되면 아무래도 특정 클래스를 주로 예측하게 되고 소수 클래스에 집중하지 못하게 될 가능성이 크다.

.

그래서 클래스 별 학습 가중치를 아래와 같이 설정해주었다.

.

역수를 취해주어서 소수 클래스에 더 큰 클래스 가중치를 부여해주어 균형있게 학습이 되도록 했다.

.

베이스라인 모델은 EfficientNet-B5를 사용하였고, imagnet으로 pre-trained 된 가중치를 가져와 사용하였다. 내 경험상 EfficeintNet이 가중치수에 비해서 성능이 되게 잘 나오는 느낌이 든다.

.

그런데 학습을 돌려보니 생각보다 test acuraccy가 잘 나오지 않았다. train acc가 0.5정도, test acc가 0.3정도에서 수렴하는 모습을 보였던 것이다. 생각보다 안나와서 당황했다. 그래서 Augmentation 방법이나 Optimizer도 바꿔보고,, Loss도 바꿔보고 모델도 Effnet 말고 새로운거 써보고.. 했는데도 오히려 성능이 나빠지기만 할 뿐 성능이 너무 안나오는 것이다.

.

그런데 결국문제를 해결했고, 그것은 굉장히 사소한 문제였다. 바로 Learning Rate가 너무 크게 잡혀있던 것이다. 사실 0.0001정도면 국룰 Learning Rate가 아닌가? 하고 너무 신경을 안쓴 탓도 있다. 0.0001대신에 0.00005로 바꾸었더니 0.5에서 수렴하던 모델이 0.95정도까지 더 학습이 진행되는 모습을 볼 수 있었다.

.

역시... 모델 학습에서 제일 중요한 하이퍼파라미터는 Learning Rate라는 cs231n의 가르침을 다시 한번 깨닫게 되었다. 내가 생각하기에는 0.0001로는 더 깊게 못내려가는 골짜기를 0.00005라는 더 작은 스텝으로는 내려갈 수 있었던 모양이다.

.

최종적으로 학습을 더 완벽하게 하기 위해서 Learning Rate Scheduler를 사용해서 4 epoch마다 Learning Rate를 0.5배 해주고 학습을 진행했고, 최종적으로 train acc는 0.97, test acc는 0.80정도 나오도록 학습을 진행할 수 있었다.

.

아래는 모델의 학습 그래프이다.

.

.

학교에서 발표할 때 데모로 유튜브에 들어가서 거기에 있는 영상중에 골라보라고 하고 그 영상으로 추론을 진행했었다. 사실 잘 나와줄지 걱정했었는데 생각보다 정확하게 맞춰서 다행이었던 기억이 난다.. 아래는 그때 추론을 진행했던 사진이다.

.

.

간단한 토이프로젝트였지만 생각치도 못한 문제에서 좀 해멨던 프로젝트였던 것 같다. 재미있었고, 좋은 교훈을 얻었다.

.

Link

여기 링크로 들어가면 캐글 NoteBook을 볼 수 있다. 데이터셋은 제공되어있지 않지만, 코드를 참고하고 싶은 분들은 참고하시길 바란다.

케라스 인공지능 텐서플로 전이학습 이미지분류



Search