2022 AI 대학원 챌린지

새벽반 - Protein pretrained model(ELMO, ESM) Fine-tuning 코드 공유

공동작성자

stroke
2022.08.01 23:23 1,700 조회 language

안녕하세요. 새벽반 팀에서 진행한 Pretrained model(ELMO, ESM)을 활용하여 Fine-tuning 하는 코드를 공유합니다.

저희 팀은 대회 초기에 시계열 데이터(epitope_seq, antigen_seq)를 중심으로 데이터 분석 및 모델 개발을 해보았습니다.
여러 RNN, GRU, LSTM 등 여러 Sequential model 등에 데이터를 적용해보았지만, 원하는 정도의 성능에 미치지 못했습니다.
이후에는 대회 운영진에서 공지한 Protein pretrained model인 ELMO(https://github.com/facebookresearch/esm) 레퍼런스를 참고하였고,
이러한 결과들을 토대로 ELMO와 ESM 모델 Fine-tuning을 진행했습니다.

epitope_seq 데이터만을 활용하여 학습하는데도 많은 GPU 메모리가 필요로 했으며 다양한 모델과 hyperparameter setting에서 실험해본 결과,
ESM_85M 모델을 Fine-tuning 했을 때의 성능이 가장 높게 나왔습니다.
공유 코드는 Training epochs 10, 20% Valid dataset으로 구성되어 있고, 이 때 F1 score가 0.58 정도 나오는 것을 확인할 수 있습니다.

코드 전체적으로는 Import libraries, Dataset 및 Data preprocessing 코드, Train, Valid, Test Dataset 및 Dataloader 선언, Train, Validation 및 Inference 코드, Pretrained model Load 및 Fine-tuning, Best model Load 및 Valid, Test data에 대한 Inference로 이뤄져 있습니다.
참가자분들이 시계열 데이터 분석을 하시는 것에 있어 유용한 정보가 되었으면 합니다.

코드
로그인이 필요합니다
0 / 1000
안녕해요
2022.08.02 10:17

안녕하세요. 코드 공유 감사합니다. ESM 부분이 ELMO부분과 같아 보이는데 혹시 확인해 주실 수 있나요?

Jin-Seop
2022.08.02 10:51

감사합니다. 똑같은 코드를 merge해서 업로드 해놓았었네요.
이를 반영해서 다시 업로드 했습니다.

안녕해요
2022.08.02 10:56

감사합니다! 좋은 코드 공유 감사합니다 ㅎㅎ