콘텐츠로 이동

PyTorch BERT-base

Overview

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

Setup & Installation

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

Note

RBLN SDK는 .whl 패키지로 배포됩니다. rebel-compiler와 rbln_driver를 사용하려면 RBLN Portal 계정이 필요합니다.

RBLN Python API 사용

모델 컴파일

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
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 런타임으로 컴파일된 모델을 로드하여 추론을 실행한 뒤, 마스킹된 토큰에 대한 예측 단어를 출력합니다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
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}))

결과는 다음과 같이 표시됩니다: 

1
[{'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로 입력을 준비합니다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
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이 필요합니다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
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)}')

결과는 다음과 같이 표시됩니다: 

1
Predicted word: white

Summary and References

이 튜토리얼에서는 RBLN Python API와 PyTorch의 torch.compile() API를 사용해 Hugging Face BERT-base 모델을 컴파일하고 추론하는 방법을 보여주었습니다.
컴파일된 모델은 RBLN NPU에서 효율적으로 마스킹 언어 모델링 태스크를 수행할 수 있습니다.

References: