최근에는 pytorch 사용자도 많이 늘고 소스코드도 많이 최적화되서 pytorch와 tensorflow를 비교해보면 좋을 듯 해서 pytorch와 tensorflow에서 BERT 인퍼런스 속도를 비교해보려고 합니다.
먼저 pytorch와 tensorflow는 둘 다 딥러닝 프레임워크로, tensorflow는 구글에서 pytorch는 페이스 북에서 개발하고 있습니다. 제가 처음 딥러닝 공부를 시작했던 2016년에는 tensorflow가 많이 압도적이였는데, 학계에서 많이 활용되면서 최근엔 pytorch 사용자가 많이 증가하고 있습니다.
자연어처리에서 BERT는 모르면 간첩이라고 할 수 있을 정도로 유명한 모델이죠. 구글에서 2018년 10월에 공개한 Transformer 기반의 거대한 언어모델입니다. Transformer의 encoder만 활용하여 만든 거대한 언어모델로 대다수의 자연어처리 솔루션에서 최신 성능을 내어 혜성처럼 등장한 모델입니다.
그래서 최근에는 BERT를 기반으로 하는 자연어처리 솔루션도 많이 등장하고 있는데, 거대한 언어모델이라 인퍼런스 속도 때문에 조금 난감해 하시는 분들이 계실거 같아서, 프레임워크만 변경해도 속도가 빨라진다면 프레임워크라도 변경을 해보시는게 도움이 되는 경우가 생길까 하여 이런 비교를 해보게 되었습니다.
결론부터 말씀드리면 pytorch가 빠릅니다. 저도 원래는 tensorflow를 많이 사용했었는데, 최근 pytorch가 장족의 발전을 거듭하여 많이 빨라졌다는 말이 많아서 궁금하여 테스트를 해봤는데, 많이 빠르네요.
소스코드와 모델은 huggingface의 transformers를 활용하였고, input에 대하여 BERT의 맨 마지막 층 vector를 계산하는 것을 기준으로 확인하였습니다.
GPU는 2060 super 기준이며, BERT base 모델의 경우 최대 512 토큰을 사용할 수 있으므로 512 token에 배치는 1개 기준으로 1000번 인퍼런스하는 속도를 측정했으며 속도 측정 결과는 아래 표와 같습니다.
| pytorch | tensorflow |
| 23.22 sec | 69.51 sec |
pytorch에서는 23.22 초가 걸렸고 tensorflow에서는 69.51 초가 걸렸습니다.
그러면 1번 기준으로는 23.22 ms, 69.51 ms가 걸린다고 생각하면 되겠죠?
테스트에 활용한 소스코드는 아래에 첨부하였습니다.
보시다시피 huggingface의 transformers를 활용하면 BERT를 쉽게 활용할 수 있습니다.
# pytorch test
# torch == 1.4.0
import time
import torch
from transformers import BertTokenizer, BertModel
device = torch.device('cuda')
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')
model.to(device)
temp = ''
for _ in range(510):
temp += 'hello '
input_ids = torch.tensor(tokenizer.encode(temp, add_special_tokens=True), device=device).unsqueeze(0)
outputs = model(input_ids)
last_hidden_states = outputs[0]
start = time.time()
for i in range(1000):
outputs = model(input_ids)
end = time.time()
print(end - start)
# tensorflow test
# tensorflow-gpu == 2.1.0
import time
import tensorflow as tf
from transformers import BertTokenizer, TFBertModel
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = TFBertModel.from_pretrained('bert-base-uncased')
temp = ''
for _ in range(510):
temp += 'hello '
input_ids = tf.constant(tokenizer.encode(temp))[None, :]
outputs = model(input_ids)
last_hidden_states = outputs[0]
start = time.time()
for i in range(1000):
outputs = model(input_ids)
end = time.time()
print(end - start)
pytorch가 빠르더라도 아주 조금 더 빠르지 않을까 생각했는데, 2배 이상 빠른 결과가 나와서 매우 놀랐습니다.
회사에서 BERT 기반으로 솔루션을 많이 개발 중인데, 이제부턴 pytorch를 활용하도록 해봐야겠습니다.
BERT 뿐만 아니라, 다른 모델들도 비교해보면 재밌을 것 같네요.