허깅페이스(Huggingface) 모델 inference(pipeline) GPU로 돌리기

 

허깅페이스를 사용하다보면 이미 만들어진 모델을 불러와서 inference만 하고자 하는 때가 있다. 가령 이미 학습된 NER 모델을 불러와 데이터 처리를 한다거나, Summarization 모델을 불러와서 필요한 문서에 대한 summarization을 수행하려고 하는 경우가 그러하다. 보통 이런 때엔 허깅페이스의 pipeline 기능을 사용한다. 그런데 처리하고자 하는 데이터가 대용량인 경우 시간이 오래 걸리는 경우가 많다. 이런 때 GPU를 사용해 model infernece를 돌리는 방법에 대해 소개한다.

참고로 이 방법은 실험을 하면서 내가 필요했던 기능인데 구글 검색을 통해 찾아봐도 나오지 않아서 ...ㅜㅜ 혼자 공식문서를 읽어서 파악하고 직접 실험을 돌리며 알아낸 것이다....!

from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline

tokenizer = AutoTokenizer.from_pretrained("dslim/bert-large-NER")
model = AutoModelForTokenClassification.from_pretrained("dslim/bert-large-NER")
ner = pipeline("ner", model=model, tokenizer=tokenizer, device=1)
print(ner(large_data))

위 예시는 "dslim/bert-large-NER"이라는 NER task에 pre-train된 bert-large모델을 사용해 large_text라는 데이터에 NER을 수행하는 코드이다. 자신이 처리하려는 task에 따라 적절히 모델이름을 바꾸거나 task 이름을 바꾸면 된다.

위와 같이 pipeline에 device 파라미터를 설정해주면 된다. 공식문서에 따르면 기본값은 -1로 설정되어있는데 이는 CPU를 사용해 model inference(모델 추론)을 돌린다는 것이다. 이를 사용하고자 하는 GPU의 넘버로 값을 바꾸어 전달해주면 해당 GPU를 사용해 추론을 돌릴 수 있다!!