KcBERT Finetune with PyTorch-Lightning v1.3.0

KcBERT Finetune with PyTorch-Lightning v1.3.0

들어가며

KcBERT를 공개하며 NSMC를 예제로 PyTorch-Lightning을 이용한 Downstream task Fine-tune을 진행하는 Colab 예제(링크, 새 창)를 만들어 배포해보았다.

한편, Transformers의 버전과 PyTorch-Lightning, 그리고 PyTorch의 버전 자체가 올라가면서 여러가지의 기능이 추가되고, 여러 함수나 내장 세팅 등이 꽤나 많이 Deprecated되었다.

사람은 언제나 귀찮음에 지배되기 때문에 코드를 한번 만들고 최소한의 수정만을 하면서 ‘돌아가기는 하는’ 수준으로 코드를 유지했다.

하지만 실행시마다 뜨는 ‘Deprecate warnings’에 질리는 순간이 오고, 지금이 바로 그 순간이라 기존 코드를 “보다 좋게”, 그리고 기왕 수정하는 김에 모델 체크포인트 저장(성능에 따른) / Logging 서비스 연동 / Inference 코드 추가를 진행해보면 어떨까 싶다.

이 글은 이 Colab 노트북(링크) 에서 직접 실행해 보실 수 있습니다.

Args? Hparams?

AS-IS

우선 가장 이슈였던 부분이자 모델 저장이 되지 않는 만악의 근원은 바로 class Args 였다.

원래는 Tap 이라는 라이브러리를 이용해 CLI에서 args를 쉽게 오버라이딩 할 수 있도록 만드는 것이 목적이었지만… 예제 코드를 Google Colab에서 실행할 수 있도록 제작을 변경하다보니, 단순한 dataset class로 대체해 설정을 단순화하는 방향으로 진행했다.

아래와 같이 값을 넣을 수 있도록 되어있고, 각 값은 args.random_seed 와 같이 액세스 하거나 오버라이딩 할 수 있도록 세팅을 해 두었다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class Arg:
random_seed: int = 42 # Random Seed
pretrained_model: str = 'beomi/kcbert-large' # Transformers PLM name
pretrained_tokenizer: str = '' # Optional, Transformers Tokenizer Name. Overrides `pretrained_model`
auto_batch_size: str = 'power' # Let PyTorch Lightening find the best batch size
batch_size: int = 0 # Optional, Train/Eval Batch Size. Overrides `auto_batch_size`
lr: float = 5e-6 # Starting Learning Rate
epochs: int = 20 # Max Epochs
max_length: int = 150 # Max Length input size
report_cycle: int = 100 # Report (Train Metrics) Cycle
train_data_path: str = "nsmc/ratings_train.txt" # Train Dataset file
val_data_path: str = "nsmc/ratings_test.txt" # Validation Dataset file
cpu_workers: int = os.cpu_count() # Multi cpu workers
test_mode: bool = False # Test Mode enables `fast_dev_run`
optimizer: str = 'AdamW' # AdamW vs AdamP
lr_scheduler: str = 'exp' # ExponentialLR vs CosineAnnealingWarmRestarts
fp16: bool = False # Enable train on FP16
tpu_cores: int = 0 # Enable TPU with 1 core or 8 cores

args = Arg()

이 부분까지는 크게 문제가 없어보인다. 나름 주석도 괜찮게 되었지 않나? 🤣

하지만 이 args 는 아주 심각한 문제가 있다. 바로 PyTorch Lightning에서 지원하는 hparams 속성에서 json으로 serializable하지 않다는 것. 이 점이 save된 ckpt에 hparams가 {} 으로 공백으로 비어있는 상황이 발생한다.

TO-BE

JSON Serializable하게 바꾸면 된다. 즉, Python dict로 바꿔주자.

(그리고 쓸모없는 인자들도 좀 지워주자…)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
args = {
'random_seed': 42, # Random Seed
'pretrained_model': 'beomi/kcbert-base', # Transformers PLM name
'pretrained_tokenizer': '', # Optional, Transformers Tokenizer Name. Overrides `pretrained_model`
'batch_size': 32,
'lr': 5e-6, # Starting Learning Rate
'epochs': 20, # Max Epochs
'max_length': 150, # Max Length input size
'train_data_path': "nsmc/ratings_train.txt", # Train Dataset file
'val_data_path': "nsmc/ratings_test.txt", # Validation Dataset file
'test_mode': False, # Test Mode enables `fast_dev_run`
'optimizer': 'AdamW', # AdamW vs AdamP
'lr_scheduler': 'exp', # ExponentialLR vs CosineAnnealingWarmRestarts
'fp16': True, # Enable train on FP16
'tpu_cores': 0, # Enable TPU with 1 core or 8 cores
'cpu_workers': 4,
}

물론 이제는 args.batch_size 로 액세스하는 것은 불가능하다. 하지만 그 대신 hparams로 바꿀 수 있다.

PyTorch-Lightning의 hparams

PyTorch-Lightning에서의 HyperParams는 아래와 같은 형식으로 세팅해 줄 경우, 자동으로 model.harpams 내에 저장된다.

1
2
3
4
5
6
7
8
9
10
11
12
>>> class ManuallyArgsModel(LightningModule):
... def __init__(self, arg1, arg2, arg3):
... super().__init__()
... # manually assign arguments
... self.save_hyperparameters()
... def forward(self, *args, **kwargs):
... ...
>>> model = ManuallyArgsModel(1, 'abc', 3.14)
>>> model.hparams
"arg1": 1
"arg2": 'abc'
"arg3": 3.14

Reference: https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#save-hyperparameters

하지만 모든 값을 모델 입력에 직접 넣어주는 것은 상당히 귀찮다.

따라서 **kwargs 를 사용해 모델 Initialize에 넣어준 뒤 self.save_hyperparameters() 를 실행해주면, 해당 값들을 모두 self.hparams 에서 액세스 할 수 있다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class Model(LightningModule):
def __init__(self, **kwargs):
super().__init__()
self.save_hyperparameters() # 이 부분에서 self.hparams에 위 kwargs가 저장된다.

self.bert = BertForSequenceClassification.from_pretrained(self.hparams.pretrained_model)
self.tokenizer = BertTokenizer.from_pretrained(
self.hparams.pretrained_tokenizer
if self.hparams.pretrained_tokenizer
else self.hparams.pretrained_model
)

[... 중략 ...]

model = Model(**args)

Logging

AS-IS

현재는 validation_epoch_end 부분에서 Tensorboard에 로깅할 값들을 아래와 같이 Dict 중 log 라는 Key의 value로 새로운 dict를 전달해주는 방식으로 로깅이 이루어지고 있다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def validation_epoch_end(self, outputs):
loss = torch.tensor(0, dtype=torch.float)
for i in outputs:
loss += i['loss'].cpu().detach()
_loss = loss / len(outputs)

loss = float(_loss)
y_true = []
y_pred = []

for i in outputs:
y_true += i['y_true']
y_pred += i['y_pred']

# Acc, Precision, Recall, F1
metrics = [
metric(y_true=y_true, y_pred=y_pred)
for metric in
(accuracy_score, precision_score, recall_score, f1_score)
]

tensorboard_logs = {
'val_loss': loss,
'val_acc': metrics[0],
'val_precision': metrics[1],
'val_recall': metrics[2],
'val_f1': metrics[3],
}

print()
pprint(tensorboard_logs)
return {'loss': _loss, 'log': tensorboard_logs}

하지만 현재의 방식은, 아래와 같이 PyTorch-Lightning v0.9.1 에서 Deprecated 되었고 v1.0.0 에 Remove 될 것이라고 한다.

(실제로 1.0.0 버전에서 사라지지는 않았고, 여전히 지원하고 있기는 하다.)

1
2
3
4
5
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/distributed.py:50: UserWarning: The {log:dict keyword} was deprecated in 0.9.1 and will be removed in 1.0.0
Please use self.log(...) inside the lightningModule instead.
# log on a step or aggregate epoch metric to the logger and/or progress bar (inside LightningModule)
self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)
warnings.warn(*args, **kwargs)

TO-BE

따라서 위에서 제공하는 것과 같이 self.log() 기능을 이용해야 한다.

1
self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)

실제 코드로 변경한 것은 아래 중복코드 제거와 함께 진행해보았다.

중복되는 코드 제거하기

AS-IS

기존의 코드는 training_step, validation_step, validation_epoch_end 세 가지로 쪼개져있는데, 실제로는 step 부분이 완전히 동일하기 때문에 이렇게 중복으로 작성할 이유가 전혀 없다.

또한, training step에서 train loss등을 출력하면, Train dataset에 대해 전체적인 metric이 나오는 것이 아니기 때문에 제거하는 것이 나았다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
def training_step(self, batch, batch_idx):
data, labels = batch
output = self(input_ids=data, labels=labels)

# Transformers 4.0.0+
loss = output.loss
logits = output.logits

preds = logits.argmax(dim=-1)

y_true = labels.cpu().numpy()
y_pred = preds.cpu().numpy()

# Acc, Precision, Recall, F1
metrics = [
metric(y_true=y_true, y_pred=y_pred)
for metric in
(accuracy_score, precision_score, recall_score, f1_score)
]

tensorboard_logs = {
'train_loss': loss.cpu().detach().numpy().tolist(),
'train_acc': metrics[0],
'train_precision': metrics[1],
'train_recall': metrics[2],
'train_f1': metrics[3],
}
if (batch_idx % self.args.report_cycle) == 0:
print()
pprint(tensorboard_logs)
return {'loss': loss, 'log': tensorboard_logs}

def validation_step(self, batch, batch_idx):
data, labels = batch
output = self(input_ids=data, labels=labels)

# Transformers 4.0.0+
loss = output.loss
logits = output.logits

preds = logits.argmax(dim=-1)

y_true = list(labels.cpu().numpy())
y_pred = list(preds.cpu().numpy())

return {
'loss': loss,
'y_true': y_true,
'y_pred': y_pred,
}

TO-BE

따라서 아래와 같이 공통함수 step 을 만들어서 training_stepvalidation_step 각각에 대해 동일하게 값을 return하도록 만들어주면 된다.

또한 step 내에서 logging을 하는 것을 제거해준다.

그리고 해당 Logging을 각 epoch end에서 처리해주기 위해

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def step(self, batch, batch_idx):
data, labels = batch
output = self(input_ids=data, labels=labels)

# Transformers 4.0.0+
loss = output.loss
logits = output.logits

preds = logits.argmax(dim=-1)

y_true = list(labels.cpu().numpy())
y_pred = list(preds.cpu().numpy())

return {
'loss': loss,
'y_true': y_true,
'y_pred': y_pred,
}

def training_step(self, batch, batch_idx):
return self.step(batch, batch_idx)

def validation_step(self, batch, batch_idx):
return self.step(batch, batch_idx)

def epoch_end(self, outputs, state='train'):
loss = torch.tensor(0, dtype=torch.float)
for i in outputs:
loss += i['loss'].cpu().detach()
loss = loss / len(outputs)

y_true = []
y_pred = []
for i in outputs:
y_true += i['y_true']
y_pred += i['y_pred']

self.log(state+'_loss', float(loss), on_step=True, on_epoch=True, prog_bar=True)
self.log(state+'_acc', accuracy_score(y_true, y_pred), on_step=True, on_epoch=True, prog_bar=True)
self.log(state+'_precision', precision_score(y_true, y_pred), on_step=True, on_epoch=True, prog_bar=True)
self.log(state+'_recall', recall_score(y_true, y_pred), on_step=True, on_epoch=True, prog_bar=True)
self.log(state+'_f1', f1_score(y_true, y_pred), on_step=True, on_epoch=True, prog_bar=True)
return {'loss': loss}

def train_epoch_end(self, outputs):
return self.epoch_end(outputs, state='train')

def validation_epoch_end(self, outputs):
self.epoch_end(outputs, state='val') # validation_epoch_end는 아무것도 반환하지 않아야 함.

위처럼 바꿔주면 학습시 아래와 같이 로그가 남는다.

1
2620/6251 [21:19<29:32, 2.05it/s, loss=0.273, v_num=0, val_loss=0.266, val_acc=0.887, val_precision=0.914, val_recall=0.856, val_f1=0.884]

validation_epoch_end 함수는 아무것도 반환하지 않아야 한다. 만약 어떤 값을 반환한다면 아래와 같은 warning이 뜬다. (Train에서야 Loss를 반환받아서 Gradient + backprop을 수행해야 하지만 Validation은 Forward만 계산하는거라 굳이 Loss를 반환할 필요가 없기 때문.)

1
UserWarning: The validation_epoch_end should not return anything as of 9.1. To log, use self.log(...) or self.write(...) directly in the LightningModule

DataLoader 코드 중복 제거

AS-IS

DataLoader를 커스텀으로 정의할 때 path, shuffle 설정을 제외하고서는 모두 동일한 코드다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
def train_dataloader(self):
df = self.read_data(self.args.train_data_path)
df = self.preprocess_dataframe(df)

dataset = TensorDataset(
torch.tensor(df['document'].to_list(), dtype=torch.long),
torch.tensor(df['label'].to_list(), dtype=torch.long),
)
return DataLoader(
dataset,
batch_size=self.args.batch_size or self.batch_size,
shuffle=True,
num_workers=self.args.cpu_workers,
)

def val_dataloader(self):
df = self.read_data(self.args.val_data_path)
df = self.preprocess_dataframe(df)

dataset = TensorDataset(
torch.tensor(df['document'].to_list(), dtype=torch.long),
torch.tensor(df['label'].to_list(), dtype=torch.long),
)
return DataLoader(
dataset,
batch_size=self.args.batch_size or self.batch_size,
shuffle=False,
num_workers=self.args.cpu_workers,
)

TO-BE

따라서 아래와 같이 dataloader 공용 함수를 만들어 정리를 해주었다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
def dataloader(self, path, shuffle=False):
df = self.read_data(path)
df = self.preprocess_dataframe(df)

dataset = TensorDataset(
torch.tensor(df['document'].to_list(), dtype=torch.long),
torch.tensor(df['label'].to_list(), dtype=torch.long),
)
return DataLoader(
dataset,
batch_size=self.hparams.batch_size or self.batch_size,
shuffle=shuffle,
num_workers=self.hparams.cpu_workers,
)

def train_dataloader(self):
return self.dataloader(self.hparams.train_data_path, shuffle=True)

def val_dataloader(self):
return self.dataloader(self.hparams.val_data_path, shuffle=False)

Model CKPT Callback

AS-IS

현재는 작업폴더(ipynb 있는 곳) 내의 ./lightning_logs/version_5 형식과 같은 폴더에 Tensorboard형식의 로그, 그리고 마지막 step의 체크포인트가 저장된다.

하지만 마지막 step이 아닌, logging에서 나타나는 여러 값(val_acc 등)을 모니터링하며 가장 성능이 좋은 모델을 찾기 위해서는 추가적인 작업이 필요하다.

TO-BE

아래와 같이 ModelCheckpoint 콜백을 사용해 아래와 같이 어떤 파일 이름으로, 어떤 Metric을 모니터링할지, 몇 개의 파일을 저장할지 지정해주면 결과값과 함께 데이터를 저장할 수 있다.

아래 코드에서는 epoch, val_acc 기준 소수점 아래 4자리까지 기록해 저장한다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
from pytorch_lightning.callbacks import ModelCheckpoint

checkpoint_callback = ModelCheckpoint(
filename='epoch{epoch}-val_acc{val_acc:.4f}',
monitor='val_acc',
save_top_k=3,
mode='max',
auto_insert_metric_name=False,
)

trainer = Trainer(
callbacks=[checkpoint_callback],
...
)

NOTE: 글 작성 시기인 2021.03.16 기준 pytorch-lightning의 Stable버전이 v1.2.3으로, auto_insert_metric_name 옵션을 지원하지 않습니다. 따라서 Colab code에서는 master branch를 받아 설치합니다.

Main Function (for DDP) 수정

AS-IS

1
2
3
4
5
def main():
print("Using PyTorch Ver", torch.__version__)
print("Fix Seed:", args.random_seed)
seed_everything(args.random_seed)
model = Model(args)

TO-BE

위에서 사용한 args 인자에 attribute로 액세스 하지 못하기 때문에 아래처럼 전달해줘야 한다.

또한, Model Init 인자에 Dict를 kwargs 형식으로 전달해주기 위해 아래와 같이 dic을 unpacking해서 전달해준다.

1
2
3
4
5
def main():
print("Using PyTorch Ver", torch.__version__)
print("Fix Seed:", args['random_seed'])
seed_everything(args['random_seed'])
model = Model(**args)

맺으며

이전 게시글, BertForSequenceClassification on Transformers v4.0.0을 맺으며 말한 것과 동일하게 마무리한다.

“잘 돌아가는 코드는 건드리는게 아니지만, 코드는 가만 내비두면 썩는다.”

Your browser is out-of-date!

Update your browser to view this website correctly. Update my browser now

×