[PyTorch] Embedding 추가하기

hugginface의 여러 모델들을 사용하다보면 vocab_size가 달라 임베딩 차원이 안맞는 경우가 생긴다. 이때 임베딩 차원을 추가하면 문제를 해결할 수 있다.

 

1. Embedding 추가하기

torch.nn.Embedding 모델에 새로운 아이템들을 더할 수 있는 방법은 아래와 같다.

vocab_size = 10
emb_dim = 128

# 기존에 학습 완료된 임베딩
original_emb = torch.nn.Embedding(vocab_size, emb_dim)

# 추가하거나 새로 학습하고자 하는 vocab의 임베딩
new_vocab_size = 2
to_add_emb = torch.nn.Embedding(new_vocab_size, emb_dim)


# 각 임베딩의 parameter를 합쳐 임베딩 차원 늘리기
new_emb_weights = torch.cat((original_emb.weight, to_add_emb.weight), dim=0)
new_emb = original_emb.from_pretrained(new_emb_weights)

print(new_emb)
>>>> Embedding(12, 128)

 

 

[참고링크]

• stackoverflow