콘텐츠로 이동

PyTorch BERT-base

개요

이 튜토리얼에서는 HuggingFace BERT-base 모델을 RBLN Python API와 PyTorch의 torch.compile() API를 사용하여 컴파일하고 추론을 실행하는 방법을 보여줍니다.

환경 설정 및 설치 확인

시작하기 전에 시스템 환경이 올바르게 구성되어 있으며, 필요한 모든 필수 패키지가 설치되어 있는지 확인하십시오. 다음 항목이 포함됩니다:

Note

  • rebel-compiler를 사용하려면 RBLN Portal 계정이 필요합니다.
  • 위 명령은 Debian 계열 Linux(예: Ubuntu)에서 pip로 패키지를 설치하는 일반적인 절차를 전제로 합니다. OS나 환경이 다른 경우에는 설치 가이드에서 지원되는 설치 조합과 적용 가능한 명령을 확인하세요.

RBLN Python API 사용

모델 컴파일

Hugging Face에서 사전 학습된 BERT-base 모델을 불러와 추론 모드로 설정한 후, RBLN 컴파일러로 컴파일하고 로컬 스토리지에 저장합니다.

import torch  
from transformers import BertForMaskedLM  
import rebel  # RBLN Compiler  

# Instantiate the BERT-base model  
bert_model = BertForMaskedLM.from_pretrained('bert-base-uncased', return_dict=False)  
bert_model.eval()  

# Compile the model  
MAX_SEQ_LEN = 128  
input_info = [  
    ('input_ids', [1, MAX_SEQ_LEN], 'int64'),  
    ('attention_mask', [1, MAX_SEQ_LEN], 'int64'),  
    ('token_type_ids', [1, MAX_SEQ_LEN], 'int64'),  
]  
compiled_model = rebel.compile_from_torch(  
    bert_model,  
    input_info,  
)  
compiled_model.save('bert_base.rbln')  

모델 추론 및 결과 확인

BertTokenizer로 입력을 준비하고, RBLN 런타임으로 컴파일된 모델을 로드하여 추론을 실행한 뒤, 마스킹된 토큰에 대한 예측 단어를 출력합니다.

import torch  
from transformers import BertTokenizer, pipeline  
import rebel  # RBLN Runtime  

MAX_SEQ_LEN = 128  
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')  
text = 'the color of rose is [MASK].'  
inputs = tokenizer(text, return_tensors='pt', padding='max_length', max_length=MAX_SEQ_LEN)  

# Load the compiled model  
module = rebel.Runtime('bert_base.rbln', tensor_type='pt')  

# Run inference  
output = module(**inputs)  

# Display results  
unmasker = pipeline('fill-mask', model='bert-base-uncased', framework='pt')  
print(unmasker.postprocess({'input_ids': inputs.input_ids, 'logits': output}))  

예시 출력:

[{'score': 0.23419028520584106, 'token': 2317, 'token_str': 'white', 'sequence': 'the color of rose is white.'}, {'score': 0.1072201207280159, 'token': 2417, 'token_str': 'red', 'sequence': 'the color of rose is red.'}, {'score': 0.07844392210245132, 'token': 2304, 'token_str': 'black', 'sequence': 'the color of rose is black.'}, {'score': 0.07031667977571487, 'token': 3756, 'token_str': 'yellow', 'sequence': 'the color of rose is yellow.'}, {'score': 0.051444780081510544, 'token': 2630, 'token_str': 'blue', 'sequence': 'the color of rose is blue.'}]

torch.compile() API 사용

모델 및 입력 준비

BERT-base 모델을 불러오고 BertTokenizer로 입력을 준비합니다.

import torch  
from transformers import BertForMaskedLM, BertTokenizer  

if torch.__version__ >= '2.5.0':  
    torch._dynamo.config.inline_inbuilt_nn_modules = False  

# Instantiate the model  
bert_model = BertForMaskedLM.from_pretrained('bert-base-uncased')  
bert_model.eval()  

MAX_SEQ_LEN = 128  
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')  
text = 'The color of a rose is [MASK].'  
inputs = tokenizer(text, return_tensors='pt', padding='max_length', max_length=MAX_SEQ_LEN)  

컴파일 및 추론

torch.compile()에서 backend='rbln'을 지정하여 RBLN 백엔드를 사용해 모델을 컴파일하고, 첫 번째 포워드 패스로 컴파일을 트리거한 후 입력에 대해 추론을 실행합니다.

Note

torch.compile()에서 RBLN 백엔드를 사용하려면 backend='rbln'을 지정해야 하며, 컴파일 전에 import rebel이 필요합니다.

import rebel  # Need to use torch dynamo's "rbln" backend.

compiled_model = torch.compile(  
    bert_model,  
    backend='rbln',  
    options={'cache_dir': './.cache'},  
    dynamic=False  
)  

# Trigger compilation with the first forward pass  
compiled_model(**inputs)  

# Run inference  
logits = compiled_model(**inputs).logits  

# Get the mask token index  
mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]  

# Get the predicted token ID  
predicted_token_id = logits[0, mask_token_index].argmax(axis=-1)  

# Print the predicted word  
print(f'Predicted word: {tokenizer.decode(predicted_token_id)}')  

예시 출력:

Predicted word: white

참고