내가 만든 ELECTRA를 Huggingface Transformers로 Porting하기

내가 만든 ELECTRA를 Huggingface Transformers로 Porting하기

BERT, ALBERT, ELECTRA 등을 직접 Pretrain하게 되면 모델이 Tensorflow의 ckpt 형태로 저장이 된다.

이번 글에서는 tensorflow ckpt를 transformers의 pytorch ckpt로 변환하는 법을 알아보겠다🤗

Intro

  • 이번 글에서는 ELECTRA-Small을 기준으로 실습을 해본다. (BERT 등도 방법은 크게 다르지 않다)
  • transformers v2.8.0을 기준으로 작성하였다. 이후 버전에서 호환되지 않는 경우가 있을 수 있다.
  • TransformersELECTRAdiscriminatorgenerator각각 따로 만들어줘야 한다!

Prerequisite

1. Original Tensorflow Checkpoint

당연히 Tensorflow로 학습한 결과물을 가지고 있어야 한다.

.
├── koelectra-small-tf
│ ├── checkpoint
│ ├── events.out.tfevents.1586942968.koelectra-small
│ ├── graph.pbtxt
│ ├── ...
│ ├── model.ckpt-700000.data-00000-of-00001
│ ├── model.ckpt-700000.index
│ └── model.ckpt-700000.meta
└── ...
checkpoint
model_checkpoint_path: "model.ckpt-700000"
all_model_checkpoint_paths: "model.ckpt-600000"
all_model_checkpoint_paths: "model.ckpt-625000"
all_model_checkpoint_paths: "model.ckpt-650000"
all_model_checkpoint_paths: "model.ckpt-675000"
all_model_checkpoint_paths: "model.ckpt-700000"

주의할 점은 checkpoint 파일에서 model_checkpoint_path 값을 “원하는 step의 ckpt”로 바꿔줘야 한다는 것이다.

2. config.json

(주의!) transformers 라이브러리가 업데이트되면서 API가 변경되는 경우가 있고, 이에 따라 config.json의 attribute가 추가/변경되는 경우가 있다.

https://huggingface.co/models로 가서 대표 모델의 config.json을 보면서 직접 만들어야 한다.

  • vocab_size 변경에만 주의하면 충분함
  • 만일 max_seq_length를 바꿨다면 max_position_embeddings도 바꿔야 함
discriminator
{
"architectures": ["ElectraForPreTraining"],
"attention_probs_dropout_prob": 0.1,
"embedding_size": 128,
"hidden_act": "gelu",
"hidden_dropout_prob": 0.1,
"hidden_size": 256,
"initializer_range": 0.02,
"intermediate_size": 1024,
"layer_norm_eps": 1e-12,
"max_position_embeddings": 512,
"model_type": "electra",
"num_attention_heads": 4,
"num_hidden_layers": 12,
"pad_token_id": 0,
"type_vocab_size": 2,
"vocab_size": 32200
}
generator
{
"architectures": ["ElectraForMaskedLM"],
"attention_probs_dropout_prob": 0.1,
"embedding_size": 128,
"hidden_act": "gelu",
"hidden_dropout_prob": 0.1,
"hidden_size": 256,
"initializer_range": 0.02,
"intermediate_size": 1024,
"layer_norm_eps": 1e-12,
"max_position_embeddings": 512,
"model_type": "electra",
"num_attention_heads": 4,
"num_hidden_layers": 12,
"pad_token_id": 0,
"type_vocab_size": 2,
"vocab_size": 32200
}

3. tokenizer_config.json

cased 모델의 경우 그냥 tokenizer를 load하면 매번 do_lower_case=False를 직접 추가해줘야 한다.

from transformers import ElectraTokenizer

tokenizer = ElectraTokenizer.from_pretrained("monologg/koelectra-small-discriminator",
do_lower_case=False)

tokenizer_config.json을 만들어주면 이러한 번거로움을 없앨 수 있다.
(만일 max_seq_length가 128이면 model_max_length도 128로 바꿔주면 된다.)

tokenizer_config.json
{
"do_lower_case": false,
"model_max_length": 512
}

4. vocab.txt

tensorflow에서 학습했을 때 쓴 vocab.txt를 그대로 쓰면 된다.

5. 최종적인 디렉토리 형태

.
├── koelectra-small-tf
│ ├── checkpoint
│ ├── events.out.tfevents.1586942968.koelectra-small
│ ├── graph.pbtxt
│ ├── ...
│ ├── model.ckpt-700000.data-00000-of-00001
│ ├── model.ckpt-700000.index
│ └── model.ckpt-700000.meta

├── electra-small-discriminator
│ ├── config.json
│ ├── tokenizer_config.json
│ └── vocab.txt

├── electra-small-generator
│ ├── config.json
│ ├── tokenizer_config.json
│ └── vocab.txt
└── ...
  • electra-small-discriminatorelectra-small-generator 폴더를 각각 만든다.
  • config.json은 discriminator용과 generator용을 따로 만들어서 폴더 안에 넣는다.
  • tokenizer_config.jsonvocab.txt는 discriminator와 generator 둘 다 동일한 파일을 넣으면 된다.

Convert

convert.py
import os
import argparse
from transformers.convert_electra_original_tf_checkpoint_to_pytorch import convert_tf_checkpoint_to_pytorch

parser = argparse.ArgumentParser()

parser.add_argument("--tf_ckpt_path", type=str, default="koelectra-small-tf")
parser.add_argument("--pt_discriminator_path", type=str, default="koelectra-small-discriminator")
parser.add_argument("--pt_generator_path", type=str, default="koelectra-small-generator")

args = parser.parse_args()

convert_tf_checkpoint_to_pytorch(tf_checkpoint_path=args.tf_ckpt_path,
config_file=os.path.join(args.pt_discriminator_path, "config.json"),
pytorch_dump_path=os.path.join(args.pt_discriminator_path, "pytorch_model.bin"),
discriminator_or_generator="discriminator")

convert_tf_checkpoint_to_pytorch(tf_checkpoint_path=args.tf_ckpt_path,
config_file=os.path.join(args.pt_generator_path, "config.json"),
pytorch_dump_path=os.path.join(args.pt_generator_path, "pytorch_model.bin"),
discriminator_or_generator="generator")

Upload your model to Huggingface s3

  1. 먼저 huggingface.co로 가서 회원가입을 해야 함
  2. 아래의 명령어로 s3에 업로드 진행
$ transformers-cli login
$ transformers-cli upload koelectra-small-discriminator
$ transformers-cli upload koelectra-small-generator

Now Let’s Use It

How to Use
from transformers import ElectraModel, ElectraTokenizer

model = ElectraModel.from_pretrained("monologg/koelectra-small-discriminator")
tokenizer = ElectraTokenizer.from_pretrained("monologg/koelectra-small-discriminator")

맺으며

생각보다 많이 사용하시네...

생각보다 많이 사용하시네...

사실 Model Porting과 관련하여 명확한 Documentation이 없어 나도 삽질을 상당히 했던 부분이다. 이번 내용이 다른 분들에게 도움이 되었으면 한다😛

또한 Huggingface s3에 모델을 업로드하는 것은 꼭 사용해보길 권한다. 이 기능이 생긴지 얼마되지 않아서 모르시는 분들이 있는데, 업로드 용량의 제한도 없고 여러모로 라이브러리 사용도 편해진다. (모델을 100개 이상 올린다고 Huggingface 팀에서 뭐라고 하지 않으니 많이들 쓰셨으면ㅎㅎ)