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

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

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

Intro#

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

Prerequisite#

1. Original Tensorflow Checkpoint#

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

1
.
2
├── koelectra-small-tf
3
│   ├── checkpoint
4
│   ├── events.out.tfevents.1586942968.koelectra-small
5
│   ├── graph.pbtxt
6
│   ├── ...
7
│   ├── model.ckpt-700000.data-00000-of-00001
8
│   ├── model.ckpt-700000.index
9
│   └── model.ckpt-700000.meta
10
└── ...

1
model_checkpoint_path: "model.ckpt-700000"
2
all_model_checkpoint_paths: "model.ckpt-600000"
3
all_model_checkpoint_paths: "model.ckpt-625000"
4
all_model_checkpoint_paths: "model.ckpt-650000"
5
all_model_checkpoint_paths: "model.ckpt-675000"
6
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도 바꿔야 함

1
{
2
  "architectures": ["ElectraForPreTraining"],
3
  "attention_probs_dropout_prob": 0.1,
4
  "embedding_size": 128,
5
  "hidden_act": "gelu",
6
  "hidden_dropout_prob": 0.1,
7
  "hidden_size": 256,
8
  "initializer_range": 0.02,
9
  "intermediate_size": 1024,
10
  "layer_norm_eps": 1e-12,
11
  "max_position_embeddings": 512,
12
  "model_type": "electra",
13
  "num_attention_heads": 4,
14
  "num_hidden_layers": 12,
15
  "pad_token_id": 0,
16
  "type_vocab_size": 2,
17
  "vocab_size": 32200
18
}

3. tokenizer_config.json#

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

1
from transformers import ElectraTokenizer
2

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

tokenizer_config.json을 만들어주면 이러한 번거로움을 없앨 수 있다.

1
{
2
  "do_lower_case": false,
3
  "model_max_length": 512
4
}

4. vocab.txt#

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

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

1
.
2
├── koelectra-small-tf
3
│   ├── checkpoint
4
│   ├── events.out.tfevents.1586942968.koelectra-small
5
│   ├── graph.pbtxt
6
│   ├── ...
7
│   ├── model.ckpt-700000.data-00000-of-00001
8
│   ├── model.ckpt-700000.index
9
│   └── model.ckpt-700000.meta
10
│
11
├── electra-small-discriminator
12
│   ├── config.json
13
│   ├── tokenizer_config.json
14
│   └── vocab.txt
15
│
16
├── electra-small-generator
17
│   ├── config.json
18
│   ├── tokenizer_config.json
19
│   └── vocab.txt
20
└── ...

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

Convert#

1
import os
2
import argparse
3
from transformers.convert_electra_original_tf_checkpoint_to_pytorch import convert_tf_checkpoint_to_pytorch
4

5
parser = argparse.ArgumentParser()
6

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

11
args = parser.parse_args()
12

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

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

Upload your model to Huggingface s3#

먼저 huggingface.co로 가서 회원가입을 해야 함
아래의 명령어로 s3에 업로드 진행

1
$ transformers-cli login
2
$ transformers-cli upload koelectra-small-discriminator
3
$ transformers-cli upload koelectra-small-generator

Now Let’s Use It#

1
from transformers import ElectraModel, ElectraTokenizer
2

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

맺으며#

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

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

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