Electra TPU Training
Electra TPU Training
build_pretraining_dataset.py에서 파일을 읽고, TFRecord 형식으로 저장되는 과정을 정리합니다.
TFRecord는 효율적인 학습을 위해 일반적인 텍스트 (.txt), 이미지(.jpg, .png) 등의 데이터를 serialize하여 binary 형식 *.tfrecord으로 저장하는 것입니다. [참고]
1. multiprocessing을 사용해서 파라미터로 받은 프로세스 개수만큼 ExampleWriter 인스턴스를 생성합니다. [link]
코드가 길어 코드에 주석을 달아놓았습니다.
def write_examples(job_id, args):
"""A single process creating and writing out pre-processed examples."""
def log(*args):
msg = " ".join(map(str, args))
print("Job {}:".format(job_id), msg)
log("Creating example writer")
example_writer = ExampleWriter( # 인스턴스 생성, 2번 참고
job_id=job_id,
vocab_file=args.vocab_file,
output_dir=args.output_dir,
max_seq_length=args.max_seq_length,
num_jobs=args.num_processes,
blanks_separate_docs=args.blanks_separate_docs,
do_lower_case=args.do_lower_case,
strip_accents=args.strip_accents,
)
log("Writing tf examples")
fnames = sorted(tf.io.gfile.listdir(args.corpus_dir))
fnames = [f for (i, f) in enumerate(fnames) # 각 프로세스가 txt파일 나눠서 가져가기
if i % args.num_processes == job_id]
random.shuffle(fnames)
start_time = time.time()
for file_no, fname in enumerate(fnames):
if file_no > 0:
elapsed = time.time() - start_time
log("processed {:}/{:} files ({:.1f}%), ELAPSED: {:}s, ETA: {:}s, "
"{:} examples written".format(
file_no, len(fnames), 100.0 * file_no / len(fnames), int(elapsed),
int((len(fnames) - file_no) / (file_no / elapsed)),
example_writer.n_written))
example_writer.write_examples(os.path.join(args.corpus_dir, fname)) # 파일 읽으러 가는 함수(참고로 파라미터는 txt파일 경로, ex) ./data/data_0.txt)
example_writer.finish()
log("Done!")
if args.num_processes == 1:
write_examples(0, args)
else:
jobs = []
for i in range(args.num_processes):
job = multiprocessing.Process(target=write_examples, args=(i, args))
jobs.append(job)
job.start()
for job in jobs:
job.join()
2. ExampleWriter init 함수
blanks_separate_docs : 빈 줄이 문서 경계를 나타내는지 여부. (default는 True이고, line by line으로 텍스트 파일을 정리했다면 False를 하면 된다. 참고로 openwebtext 데이터 예시도 False)
self._example_builder : 토큰화하고 TFRecord 형식으로 만들 친구, ExampleBuilder 파라미터로 max_seq_length가 있는데 텍스트의 길이를 더해서 max_seq_length를 넘어야 TFRecord 만들러 간다.
pretrain_data.tfrecord-0-of-* 을 미리 만들어서 리스트에 저장 (ex. pretrain_data.tfrecord-0-of-0, pretrain_data.tfrecord-0-of-1, …)
class ExampleWriter(object):
"""Writes pre-training examples to disk."""
def __init__(self, job_id, vocab_file, output_dir, max_seq_length,
num_jobs, blanks_separate_docs, do_lower_case,
num_out_files=1000, strip_accents=True):
self._blanks_separate_docs = blanks_separate_docs
tokenizer = tokenization.FullTokenizer(
vocab_file=vocab_file,
do_lower_case=do_lower_case,
strip_accents=strip_accents)
self._example_builder = ExampleBuilder(tokenizer, max_seq_length)
self._writers = []
for i in range(num_out_files):
if i % num_jobs == job_id:
output_fname = os.path.join(
output_dir, "pretrain_data.tfrecord-{:}-of-{:}".format(
i, num_out_files))
self._writers.append(tf.io.TFRecordWriter(output_fname))
self.n_written = 0
3. ExampleWriter 클래스의 write_examples 함수
참고로 write_examples 함수가 두개이다. ExampleWriter 클래스의 write_examples 함수와 1에서 적어놓은 write_examples 함수
파이썬 open()으로 작성한 test.txt 파일도 tf.gfile.GFile()을 통해 읽을 수 있으며, 그 반대도 가능하다.
–> tensorflow 용 파일 입출력 함수라 생각하면 되는 것 같다.
SerializeToString()을 통해 Serialize한 후 init에서 만들어둔 self._writers에 저장합니다.
def write_examples(self, input_file):
"""Writes out examples from the provided input file."""
with tf.io.gfile.GFile(input_file) as f:
for line in f:
line = line.strip()
if line or self._blanks_separate_docs: # 위에서 설명한 파라미터
example = self._example_builder.add_line(line)
if example:
self._writers[self.n_written % len(self._writers)].write(
example.SerializeToString())
self.n_written += 1
example = self._example_builder.add_line("") # for문이 끝나고 target_length를 못 넘어서 저장 안된 것이 있을 수 있으니까 빈문장을 하나 넣어준다.
if example: # 혹시 있었으면 위와 같이 Serialize한 후 저장.
self._writers[self.n_written % len(self._writers)].write(
example.SerializeToString())
self.n_written += 1
이렇게 하여 지정한 num_out_files 개수만큼 *.tfrecord가 만들어진다.
pretrain_data.tfrecord-0-of-0,
pretrain_data.tfrecord-0-of-1
pretrain_data.tfrecord-0-of-2,
...
pretrain_data.tfrecord-0-of-num_out_files
4. add_line 함수
3번에서 본 add_line 함수이다.
빈 문장이어도, 현재 저장되어 있는 문장이 있으면 _create_example 함수 호출
토큰화하고 레이블 값으로 바꾸어준다. 문장의 길이를 더해주고, target_length를 넘어야 _create_example 함수 호출.
def add_line(self, line):
"""Adds a line of text to the current example being built."""
line = line.strip().replace("\n", " ")
if (not line) and self._current_length != 0: # 빈 문장이어도, 현재 저장되어 있는 문장이 있으면 _create_example 함수 호출
return self._create_example()
bert_tokens = self._tokenizer.tokenize(line)
bert_tokids = self._tokenizer.convert_tokens_to_ids(bert_tokens)
self._current_sentences.append(bert_tokids)
self._current_length += len(bert_tokids)
if self._current_length >= self._target_length: # 문장의 길이를 더해주고, target_length를 넘어야 _create_example 함수 호출.
return self._create_example()
return None
5. _create_example 함수
5프로의 확률로 target_length를 5에서 self._max_length 사이의 값으로 정한다. (random.randint(5, self._max_length))
나머지는 target_length랑 _max_length랑 같다.
즉, 위에서 문장의 길이를 더한 것이 target_length를 넘어야 _create_example함수를 호출한다고 했는데 그 target_length를 정해준다.
여러 가지 조건을 통해 segment를 나눈다.
그리고 [CLS], [SEP]이 포함되지 않았기 때문에 _max_length - 2로 자르고 _make_tf_example함수 호출
def _create_example(self):
"""Creates a pre-training example from the current list of sentences."""
# small chance to only have one segment as in classification tasks
if random.random() < 0.1:
first_segment_target_length = 100000
else:
# -3 due to not yet having [CLS]/[SEP] tokens in the input text
first_segment_target_length = (self._target_length - 3) // 2
first_segment = []
second_segment = []
for sentence in self._current_sentences:
# the sentence goes to the first segment if (1) the first segment is
# empty, (2) the sentence doesn't put the first segment over length or
# (3) 50% of the time when it does put the first segment over length
if (len(first_segment) == 0 or
len(first_segment) + len(sentence) < first_segment_target_length or
(len(second_segment) == 0 and
len(first_segment) < first_segment_target_length and
random.random() < 0.5)):
first_segment += sentence
else:
second_segment += sentence
# trim to max_length while accounting for not-yet-added [CLS]/[SEP] tokens
first_segment = first_segment[:self._max_length - 2]
second_segment = second_segment[:max(0, self._max_length -
len(first_segment) - 3)]
# prepare to start building the next example
self._current_sentences = []
self._current_length = 0
# small chance for random-length instead of max_length-length example
if random.random() < 0.05: # 5프로의 확률로 target_length를 5에서 self._max_length 사이의 값으로 정한다.
self._target_length = random.randint(5, self._max_length)
else:
self._target_length = self._max_length
return self._make_tf_example(first_segment, second_segment)
6. _make_tf_example 함수
special 토큰을 추가해주고, attention mask, segment_ids 를 만들어줍니다.
tensorflow의 tf.train.Example, tf.train.Features 를 거쳐 리턴해줍니다. (이 값을 리턴 받아서 SerializeToString()을 통해 Serialize 한 후 tfrecord로 저장)
def _make_tf_example(self, first_segment, second_segment):
"""Converts two "segments" of text into a tf.train.Example."""
vocab = self._tokenizer.vocab
input_ids = [vocab["[CLS]"]] + first_segment + [vocab["[SEP]"]]
segment_ids = [0] * len(input_ids)
if second_segment:
input_ids += second_segment + [vocab["[SEP]"]]
segment_ids += [1] * (len(second_segment) + 1)
input_mask = [1] * len(input_ids)
input_ids += [0] * (self._max_length - len(input_ids))
input_mask += [0] * (self._max_length - len(input_mask))
segment_ids += [0] * (self._max_length - len(segment_ids))
tf_example = tf.train.Example(features=tf.train.Features(feature={
"input_ids": create_int_feature(input_ids),
"input_mask": create_int_feature(input_mask),
"segment_ids": create_int_feature(segment_ids)
}))
return tf_example
run_pretraining.py
tf.data.Dataset의 API를 통해 TFRecord를 읽어, 실제 학습에 이용할 수 있는 Tensor 형태로 변환한다.
각 example을 읽어서 serialize 했던 데이터들을 다시 파싱한 후 형 변환을 해줍니다. [참고1] [참고2]
따라서 구글 Official 코드를 이용하기 위해서는 WordpieceTokenizer를 사용한 vocab.txt가 필요하고, 다른 토크나이저를 사용하려면 vocab.txt와 토크나이저를 조금 구현해야한다. [고쳐야 될 부분]
Leave a comment