#인공지능

TensorFlow Progbar 사용법, 모델 저장하기 – model.fit()에서 벗어나기! (2)

TensorFlow custom trainer를 직접 구현해서 모델을 학습시킬 때, Progbar를 이용하여 학습이 얼마나 진행되었는지 표시하는 방법과 Checkpoint로 모델을 저장하는 방법을 알려드립니다.

2022-11-07 | 박정현

모두의연구소 콘텐츠 크리에이터 모임 “코크리” 1기 활동으로 작성된 글입니다. (저자 블로그)

원본 글: model.fit()에서 벗어나기! (2)
예상 독자
- model.fit()에서 벗어나고 싶은 사람
- tfds.load(‘mnist’)로 mnist 데이터 말고 새로운 데이터 혹은 커스텀 데이터를 쓰는 튜토리얼을 읽고 싶은 사람
- 커스텀 데이터를 model.fit()으로 학습하는 방법 말고 커스텀 트레이너로 학습하고 싶은 사람

model.fit()에서 벗어나기! (1) 요약

model.fit()에서 벗어나기! (1)에서 우리의 몇 가지에 대해 이야기했습니다.

tf.data를 이용한 데이터 로드하기
tf.keras.Model을 상속받아 나만의 모델 만들기
Train loop로 모델 학습하기

하지만 우리가 만들었던 간단한 train loop로는 지금 어느 정도 학습되고 있는지, loss가 어느 정도 인지를 알 수 없고, 모델을 저장해보지도 않았습니다. 오늘은 그것들을 해보려고 합니다! 추가로 argparse를 이용하여 epoch, batch 사이즈 등의 값을 읽어오도록 해보겠습니다. 불태워봅시다.

이 기능들을 추가하면서 우리는 조각 코드를 계속 보게 될 것입니다. 조각 코드가 아닌 전체의 흐름을 보고 싶다면 아래의 github에서 전체 코드를 살펴보는 것을 추천합니다!

https://github.com/parkjh688/bye_model.fit

Train loop에 기능 추가하기

1. Progbar 추가하기

Progbar란 Progression Bar의 줄임말로, 학습이 얼마나 진행되었는지 아래 [그림 2]와 같은 바 형태로 보여주는 것입니다.

이전 Trainer의 train() 함수는train_dataset, train_metric 단 두 개만을 인자로 받았습니다. 이제는 각 step 별로 Progbar를 찍어주기 위해 steps_per_epoch이라는 인자를 추가해줄 것입니다. 여기서 step이란 batch 1개의 loss를 계산하고 gradient를 한 번 업데이트하는 것을 말합니다. 그 step은 어떻게 구할 수 있을까요? 전체 학습 데이터 사이즈를 TRAIN_SIZE라고 했을 때, 그걸 batch 사이즈로 나눈 것이 step 수가 되겠죠? 코드로 바꾼다면 아래의 labmda compute_steps_per_epoch으로 계산할 수 있을 것입니다.

compute_steps_per_epoch = lambda x: int(math.ceil(1. * x / batch_size))
steps_per_epoch = compute_steps_per_epoch(TRAIN_SIZE)

하지만 우리는 아직 TRAIN_SIZE를 모릅니다. 그래서 기존의 load_data에서 데이터를 가져올 때 TRAIN_SIZE를 함께 계산하여 가져오게 할 것입니다. 바로 이렇게 말이죠!

train_ds, TRAIN_SIZE = load_data(data_path=train_path, img_shape=(224, 224), batch_size=batch_size)

그렇다면 load_data는 어떻게 바뀌어야 할까요? 아래의 코드를 보세요. DATASET_SIZE를 구하는 코드가 생겼습니다. tf.data.experimental.cardinality()로 우리는 데이터셋의 사이즈를 알 수 있습니다. 그리고 .numpy() 메소드를 호출하여 넘파이 배열로 바꾸었습니다.

def load_data(data_path, img_shape, batch_size=64, is_train=True):
    class_names = [cls for cls in os.listdir(data_path) if cls != '.DS_Store']
    data_dir = pathlib.Path(data_path)
    list_ds = tf.data.Dataset.list_files(str(data_dir / '*/*'))

    labeled_ds = list_ds.map(lambda x: process_path(x, class_names, img_shape))
    labeled_ds = prepare_for_training(labeled_ds, batch_size=batch_size, training=is_train)

    DATASET_SIZE = tf.data.experimental.cardinality(list_ds).numpy()

    return labeled_ds, DATASET_SIZE

이제 step 사이즈도 알았으니 step마다 진행 사항을 찍어줄 수 있을 것 같습니다! 진행 사항을 찍기 위해 from tensorflow.keras.utils import Progbar를 사용해보겠습니다. 우리가 원하는 Progbar의 전체 길이는 steps_per_epoch(이전에 구한 DATA_SIZE와 같음) x batch가 될 것입니다. 1 step마다 1 batch를 사용하게 되니까요. 휴, 헷갈리죠? 😵 Progbar는 아래처럼 만들면 됩니다.

metrics_names = ['train_loss']
progBar = Progbar(steps_per_epoch * self.batch, stateful_metrics=metrics_names)

그리고 매 step이 바뀔 때마다 update 함수를 사용해주면 됩니다. values에는 progbar에서 보여주고 싶은 값을 넣어주면 됩니다.

values = [('train_loss', train_loss), ('train_acc', train_acc)]
progBar.update((step_train + 1) * self.batch, values=values)

하지만 Progbar를 사용하지 않고 tqdm, 커스텀 progress bar 등 어떤 것으로 만들어도 상관없습니다.

2. Validation 데이터 추가하기

우리는 지금까지 train 데이터로만 진행을 해왔습니다. 하지만 validation 데이터가 있는 상황도 있을 것입니다. validation 데이터 추가를 했다고 해서 크게 코드에 차이가 있지 않습니다. 그래서 빠르게 validation 데이터를 추가해보도록 하겠습니다. 저는 validation 데이터가 따로 있다고 가정하고 만들었습니다. 그렇다는 건 train_path, val_path가 모두 존재한다는 뜻이겠죠! 단순히 load_data를 두 번 해주는 것에 지나치지 않습니다. 아래와 같습니다.

train_ds, TRAIN_SIZE = load_data(data_path=train_path, img_shape=(224, 224), batch_size=batch_size)
val_ds, VAL_SIZE = load_data(data_path=val_path, img_shape=(224, 224), batch_size=batch_size, is_train=False)

compute_steps_per_epoch = lambda x: int(math.ceil(1. * x / batch_size))
steps_per_epoch = compute_steps_per_epoch(TRAIN_SIZE)
val_steps = compute_steps_per_epoch(VAL_SIZE)

하지만 이 데이터를 train step에 넣어주려면 어떻게 해야 될까요? 당연히 Trainer 클래스의 train 함수에 인자를 늘려줘야겠죠. validation을 위한 Progbar 역시 만들어주도록 해봅시다. ~~오! 슬슬 복잡해지는 것처럼 보입니다. 그럴 땐 튜토리얼을 읽지 않고 github 클론이나 받아서 코드나 돌려볼까 하는 마음이 생기죠~~

for step, (x_batch_val, y_batch_val) in enumerate(val_dataset):
    logits = model(x_batch_val, training=False)
    val_loss = self.loss_fn(y_batch_val, logits)
    
    # Update val metrics
    val_acc = self.compute_acc(logits, y_batch_val)
    values = [('train_loss', train_loss), ('train_acc', train_acc), ('val_loss', val_loss), ('val_acc', val_acc)]

progBar.update((step_train + 1) * self.batch, values=values, finalize=True)

아래 [그림 3]과 같은 Progbar 결과를 얻을 수 있습니다.

사실 이 Progbar를 그리는 데에 약간의.. anger가 있었습니다. Progbar 버그 때문인데요. 하지만 오아시스 행님들도 Don’t look back in anger라고 노래하셨으니 마음을 비우며 혹시나 같은 에러가 나는 분들을 위해 제가 링크를 하나 달아놓겠습니다.

https://github.com/keras-team/keras/issues/5906

위의 링크와도 관련 있고 제가 겪기도 했던 그 버그를 살짝 이야기해드리고 싶은데요. 아래의 [그림 4]처럼 매번 뉴 라인을 찍어주는 것이 버그였습니다. 물론 버전에 따라 다를 수 있습니다. 이와 같은 버그를 겪는 분들은 케라스를 인스톨해둔 위치에 찾아가 keras/utils/generic_utils.py를 한 번 들여다보는 것을 추천합니다. 출력하는 부분에서 문제가 있을지도 모릅니다.

3. 모델 저장하기: Checkpoint

이제 학습과 관련하여 몇 가지 처리를 해놓았으니, 학습된 모델을 저장해보도록 합시다. 기쁘게도 tensorflow에는 tf.train.Checkpoint라는 것이 있습니다. 모델을 저장해주는 아주 친절한 친구입니다. 사용법은 아주 간단합니다. 쓰려고 미리 준비해뒀던 optimizer와 model을 인자로 넣어주고 checkpoint를 만든 뒤, 그 checkpoint를 checkpoint manager로 만들면 됩니다. 이때 directory에는 모델이 저장될 directory를 입력하고, max_to_keep에는 최대로 저장할 모델의 개수를 지정해줍니다.

model = YogaPose(num_classes=num_classes)
optimizer = tf.keras.optimizers.Adam(learning_rate=0.0001)

checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)
manager = tf.train.CheckpointManager(checkpoint, directory=".", max_to_keep=5)

그다음엔 역시 Trainer 클래스의 train 함수에 manager를(tf.train.CheckpointManager 객체) 보내 1 epoch이 끝날 때마다 manager.save()를 넣어 모델을 저장해주면 됩니다. 하지만 실제 tensorflow, keras에는 callback이라는 것이 있죠. 그 callback 중엔 학습한 모델 중 가장 좋은 모델만 저장해주는 애가 있습니다. 우리도 그것과 비스무리한 걸 구현해볼까 합니다.

아래를 보세요. 아! train 부분은 train_on_batch 함수로 그대로 빼뒀습니다. 그 점 참고하세요. 가장 좋은 모델을 저장하려면 임의의 loss인 best_loss를 두고, 매번 생기는 val_loss가 그것보다 낮을 때만 모델을 저장하게 하면 됩니다. loss는 0에 가까워질수록 좋으니까요. loss가 아닌 acc로 해도 무방합니다. 단 best_loss가 아닌 best_acc일 경우에는 이전의 acc보다 현재 acc가 더 높을 때 저장을 해야 한다는 것을 주의해야 합니다.

def compute_acc(self, y_pred, y):
    correct = tf.equal(tf.argmax(y_pred, 1), tf.argmax(y, 1))
    accuracy = tf.reduce_mean(tf.cast(correct, tf.float32))
    return 
        
@tf.function
def train_on_batch(self, x_batch_train, y_batch_train):
    with tf.GradientTape() as tape:
        logits = model(x_batch_train, training=True)    # 모델이 예측한 결과
        train_loss = self.loss_fn(y_batch_train, logits)     # 모델이 예측한 결과와 GT를 이용한 loss 계산

    grads = tape.gradient(train_loss, model.trainable_weights)  # gradient 계산
    self.optimizer.apply_gradients(zip(grads, model.trainable_weights))  # Otimizer에게 처리된 그라데이션 적용을 요청

    return train_loss, logits

def train(self, train_dataset, acc_metric, steps_per_epoch, val_dataset, val_step, checkpoint_manager):
    metrics_names = ['train_loss', 'train_acc', 'val_loss']

    best_loss = 100
    for epoch in range(self.epochs):
        print("\nEpoch {}/{}".format(epoch+1, self.epochs))

        train_dataset = train_dataset.shuffle(100)
        val_dataset = val_dataset.shuffle(100)

        train_dataset = train_dataset.take(steps_per_epoch)
        val_dataset = val_dataset.take(val_step)

        progBar = Progbar(steps_per_epoch * self.batch, stateful_metrics=metrics_names)

        train_loss, val_loss = 100, 100

        # 데이터 집합의 배치에 대해 반복합니다
        for step_train, (x_batch_train, y_batch_train) in enumerate(train_dataset):
            train_loss, logits = self.train_on_batch(x_batch_train, y_batch_train)

            # train metric(mean, auc, accuracy 등) 업데이트
            acc_metric.update_state(y_batch_train, logits)

            train_acc = self.compute_acc(logits, y_batch_train)
            values = [('train_loss', train_loss), ('train_acc', train_acc)]
            # print('{}'.format((step_train + 1) * self.batch))
            progBar.update((step_train + 1) * self.batch, values=values)

        for step, (x_batch_val, y_batch_val) in enumerate(val_dataset):
            logits = model(x_batch_val, training=False)
            val_loss = self.loss_fn(y_batch_val, logits)
            val_acc = self.compute_acc(logits, y_batch_val)
            values = [('train_loss', train_loss), ('train_acc', train_acc), ('val_loss', val_loss), ('val_acc', val_acc)]
        progBar.update((step_train + 1) * self.batch, values=values, finalize=True)

        if val_loss < best_loss:
            best_loss = val_loss
            print("\nSave better model")
            print(checkpoint_manager.save())

아래 [그림 5]를 보면 loss가 더 적어졌을 때는 모델을 저장하고, 클 때는 모델을 저장하지 않고 있습니다. 우리만의 작은 callback을 만든 셈이죠! 드디어 끝이 보입니다.

여기서 잠깐.

여기까지 예제 코드를 잘 만들어왔던 제가 여기서 난관에 봉착했었죠. 사실 난관은 잦았지만.. 어쨌든 Yoga Pose 이미지 데이터만으로는 모델이 생각보다 학습이 잘 안된다는 것이었습니다. 심지어 오버피팅도 안됐다고요! 그래서 제가 시도한 한 가지가 있었습니다.

간단히 설명한다면 우리가 쓰고 있는 요가 데이터는 [그림 7]의 왼쪽 이미지와 같이 생겼습니다. 학습에서의 문제는 오버피팅조차 되지 않고 모델이 학습하기 싫어한다는 것이었습니다. 어쩌면 주인을 닮았을지도 모릅니다. 저도 공부하기 싫거든요. 그래서 [그림 7]의 오른쪽 이미지같이 자세의 keypoints들을 모델에 함께 넣어주기로 했습니다. 모델에게 이미지 외에도 조금 더 많은 정보를 주는 것이죠. 그러기 위해서는 제가 가진 Yoga Pose 이미지 데이터에 있는 동작들에 대한 keypoints 값을 가지고 있어야 합니다.

이 과정은 직접 annotation 할 수도 있었지만 그 작업은 하고 싶지 않았기에, 이미 학습되어있는 MoveNet이라는 모델로부터 keypoints label을 얻는 것으로 대체하였습니다. 물론 모델을 학습하기 위해서는 정확한 label이 있는 것이 중요합니다. 하지만 갑자기 생각난 ‘keypoints를 넣어주면 모델이 좀 더 잘 학습하지 않을까?’ 라는 아이디어를 체크하기 위해서는 아주 정교한 fine label이 아닌 대충 만들어본 coarse label로도 결과를 체크해볼 수 있을 것입니다. 그 과정에서 더 많은 데이터를 가져와야 하기 때문에 load_data도 바뀌었고, 입력이 늘었으니 Model도 바뀌었습니다. 그 모든 과정은 아래의 코랩을 참고하시길 바랍니다.

Google Colaboratory

블로그 형식으로 보고 싶다면 코크리 블로그에 옮겨놓은 글로 보셔도 좋을 거예요! (내용은 같습니다)

coarse label 로 custom dataloader 만들어보기 (근데 이제 딥러닝 인퍼런스 값을 곁들인..)

4. argparse 사용하기

이제 마지막으로 argparse라는 것을 추가해보겠습니다. 없어도 되는 기능이지만 이걸 넣으면 왠지 멋있어 보이니까 그냥 하겠습니다. 멋있는 것도 중요하지만, argparse를 사용하면 매번 .py 내에서 값을 바꿔주지 않고도 python train.py 명령만으로 실험에 필요한 epoch, batch, data path 등을 설정해줄 수 있습니다! 그 동안은 아래의 코드와 같이 main 을 넣어 학습시켰습니다. 하지만 이제는 파이썬 파일 실행 시 num_classes, epoch, batch_size, data_path 등을 받을 수 있도록 바꿔보겠습니다.

if __name__ == '__main__':
    num_classes = 107
    epoch = 1000
    batch_size = 64

    train_path = '/Users/edensuperb/Downloads/pythonProject/dataset'
    val_path = '/Users/edensuperb/Downloads/pythonProject/dataset'

    train_ds, TRAIN_SIZE = load_data(data_path=train_path, img_shape=(224, 224), batch_size=batch_size)
    val_ds, VAL_SIZE = load_data(data_path=val_path, img_shape=(224, 224), batch_size=batch_size, is_train=False)

아래와 같이 만들면 될 것입니다.

import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--epoch",          type=int,       default=100)
parser.add_argument("--batch_size",     type=int,       default=64)
parser.add_argument("--num_classes",    type=int,       default=107)
parser.add_argument("--img_size",       type=int,       default=224)
parser.add_argument("--train_path",     type=str,		default='./dataset/train')
parser.add_argument("--val_path",       type=str,		default='./dataset/val')
parser.add_argument("--checkpoint_path",type=str,		default='./checkpoints')

args = parser.parse_args()

후에 받은 인자들은 args.batch_size와 같이 사용하면 됩니다. 짜잔 python train.py — batch_size 32 — checkpoint_path ‘.’ 로 실행해보았더니 원하는대로 잘 되었습니다. 위에서 batch_size가 default로 64로 지정되어있지만 제가 32로 바꿔서 시도해본 것 처럼 여러분은 batch_size를 여러분 마음대로 바꿀 수 있습니다.

5. 모델 테스트해보기

제가 잊고 있던 것이 하나 있었습니다. 바로 모델 테스트입니다. 이게 진짜 진짜 마지막입니다. 학습한 모델이 잘 테스트 되는지도 역시 중요하죠. 짧은 코드와 함께 테스트해보겠습니다. test.py라는 새 파일에 argparse로 인자를 받아보겠습니다. 방법은 간단합니다. tf.train.Checkpoint로 이미 저장된 ckpt를 가져오면 됩니다. ckpt에는 우리가 만든 모델이 그대로 저장되어 있기 때문입니다. 그 후 다시 데이터를 로드한 후 모델을 inference해서 진짜 label과 inference 결과를 비교하여, 그 값이 몇 개나 맞았는지 체크해보면 되죠. 아래는 그 과정입니다.

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--epoch",          type=int,       default=100)
    parser.add_argument("--num_classes",    type=int,       default=107)
    parser.add_argument("--img_size",       type=int,       default=224)
    parser.add_argument("--test_path",      type=str,       default='./dataset/test')
    parser.add_argument("--checkpoint_path",type=str,		default='./checkpoints/ckpt-77')

    args = parser.parse_args()

    model = YogaPose(num_classes=args.num_classes)
    optimizer = tf.keras.optimizers.Adam(learning_rate=0.0001)

    test, TEST_SIZE = load_data(data_path=args.test_path, img_shape=(args.img_size, args.img_size), batch_size=64)

    checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)
    checkpoint.restore(args.checkpoint_path)

    for step_train, (x_batch_train, y_batch_train) in enumerate(test.take(10)):
        # print(model(x_batch_train))
        prediction = model(x_batch_train)
        # print(tf.argmax(y_batch_train, axis=1))
        # print(tf.argmax(prediction, axis=1))
        # print(tf.equal(tf.argmax(y_batch_train, axis=1), tf.argmax(prediction, axis=1)))
        print("{}/{}".format(np.array(tf.equal(tf.argmax(y_batch_train, axis=1), tf.argmax(prediction, axis=1))).sum(), tf.argmax(y_batch_train, axis=1).shape[0]))
        # print("Prediction: {}".format(tf.argmax(prediction, axis=1)))

아래와 같이 나오네요. 대충 32개 중에 23개 — 28개 사이로 맞히고 있습니다.

이제 여러분들도 원하는 데이터를 로드해서 train loop로 학습하고 test까지 해볼 수 있을 것 같다는 생각이 듭니다. 글이 마음에 드셨을지 모르겠습니다. 하지만 저는 약간.. 터프 타임이었네요. 다음에 또 다른 콘텐츠로 돌아오겠습니다. 끝까지 읽어주셔서 감사드립니다.

콘텐츠 크리에이터 소개

안녕하세요! 여행, 피자, 역사를 좋아하는 머신러닝 엔지니어 박정현(Eden)입니다. 코로나를 뚫고 이탈리아(피자와 역사는 덤)를 다녀왔을 정도로 여행을 좋아합니다. 좋은 콘텐츠 많이 보여드렸으면 좋겠습니다. (블로그)

박정현

안녕하세요! 저는 박정현이라고 합니다. 현재 2D 컴퓨터 비전 분야의 머신러닝 엔지니어로 활동하고 있습니다. 최근에는 이미지 검색 쪽 일을 하고 있고, 사실 XR과는 관련이 없습니다. 하지만 개인적으로 XR에 관심이 있고 특히 포켓몬고 유저로서 AR에 가장 관심이 많습니다 ㅎㅎ 최근 반년간은 NeRF(Neural Radiance Fields)를 공부하고 있는데, 공부하면서 SLAM이나 3D reconstuction 등과 관련된 다양한 연구들에 더욱 흥미를 느끼고 있습니다. 또한, ARKit과 딥러닝 학습을 활용하여 데이터를 수집하고 모델을 개발하는 과정에도 큰 관심이 있습니다. 특히, 하드웨어에서 얻은 이미지를 활용하여 XR과 관련된 모델을 개발하고, 앱에서 상호작용이 가능하도록 활용하는 것을 해보고자 하는 마음이 있습니다. 이러한 목표를 혼자서는 큰 분야이기 때문에 지원하게 되었고, 모델을 학습시키는 방법이나 인퍼런스(추론)하는 과정은 알고 있지만 모델을 하드웨어와 상호작용이 가능하도록 사용하는 방법(HW, 앱 개발, ARKit 사용하기 등) 을 공부하고 싶어서 지원했습니다.

목록으로 돌아가기