[Tensorflow.js] 웹 브라우저를 위한 딥러닝 모델 학습

이번 게시글에서는 웹 브라우저에 Obejct detection을 적용하기 위한 딥러닝 모델 학습에 대해 다루고자 합니다.

 

현재 친구들과 함께 유튜브 영상 속 한국 음식을 인식하고 해설을 제공해주는 크롬 익스텐션을 개발하고 있습니다. 크롬 확장 프로그램에 딥러닝 모델을 적용하기 위해서는 tensorflow.js 버전의 모델이 필요합니다.

 

처음에는 yolo나 keras를 이용하여 모델을 학습시키고, 그걸 tfjs로 변환하고자 하였는데 그 과정이 생각보다 복잡하더라구요. 그래서 결국 애초부터 tensorflow를 이용하여 모델을 학습시킨 후, 이를 tfjs로 변환하여 사용하기로 결정하였습니다.

 

전체적인 프로세스는 다음과 같습니다.

1. 한식 데이터셋 제작
2. tensorflow 기반 SSD MobileNet v2 학습
3. tensorflow.js 모델로 변환

---

1. 한식 데이터셋 제작

모델을 학습시키기 위해서는 먼저 데이터셋이 필요합니다. 저희 서비스의 타겟 콘텐츠가 유튜브 영상이기 때문에, 유튜브 영상 이미지를 기반으로 직접 데이터셋을 제작하였습니다. 이때 사용한 툴이 바로 roboflow입니다.

roboflow에서 제공하는 annotation tool을 이용하면 직접 이미지를 업로드한 후 라벨링하는 작업을 보다 쉽게 수행할 수 있습니다. 하나의 프로젝트에 대한 공동 작업이 가능하기 때문에 팀원들과 함께 유용하게 사용할 수 있으며, 자동으로 train/valid/test set을 나눠준다는 것 또한 roboflow의 장점 중에 하나입니다.

 

roboflow에 로그인하여 자신만의 work space를 만든 후 이미지를 업로드하고, 간단한 드래그를 통해 다음과 같이 라벨링을 수행할 수 있습니다.

저희는 이러한 작업을 통해 한국 음식 60개에 대해 총 2만 4천장의 데이터셋을 제작하였습니다.

 

물론 roboflow에도 단점이 존재합니다.

 

(1) 수정이 어렵다.

- 이미지를 잘못 업로드한 경우, 각 이미지마다 일일이 삭제를 해줘야합니다. 그리고 클래스명이 잘못 입력한 경우에도 해당 클래스명으로 라벨링한 모든 이미지를 하나하나 수정해줘야 합니다. 일괄로 삭제하거나 수정할 수 있는 방법이 없어서 좀 불편하긴 했습니다.

 

(2) 무료 버전에서는 export할 수 있는 이미지 개수에 제한이 있다.

- 원래는 딱히 limit이 없었는데, 최근에 정책이 변경되었는지 이미지가 10,000개 이상이면 export할 때 유료 플랜으로 바꾸라고 뜨더라고요. 저희는 그전에 미리 export를 해놓아서 다행이었지, 만약 그러지 않았다면...😥 roboflow를 사용할 때 이 점 유의하세요!

---

2. Tensorflow 기반 SSD MobileNet v2 학습

딥러닝 모델을 학습하기 위해 구글 코랩을 이용하였습니다. 코랩을 이용하면 구글에서 제공하는 gpu를 이용할 수 있고, 기본적인 환경 세팅이 이미 다 되어 있기 때문에 수고스러움을 줄일 수 있습니다. 다만 사용 시간이 제한적이라는 단점이 있으니 유의하시기 바랍니다.

 

우선 모델을 학습하기 전, 필요한 파일을 설치하도록 하겠습니다.

Tensorflow Object Detection API 설치

!git clone https://github.com/tensorflow/models.git

%cd /content/models/research/
!protoc object_detection/protos/*.proto --python_out=.
# TensorFlow Object Detection API 설치
!cp object_detection/packages/tf2/setup.py .
!python -m pip install .

# modeil builder test 수행
!python /content/models/research/object_detection/builders/model_builder_tf2_test.py

Dataset 가져오기

roboflow의 장점 중에 하나가, 바로 데이터셋을 다양한 형식으로 export 할 수 있다는 점입니다. 원하는 프로젝트에서 새로운 버전을 생성한 후 export 버튼을 누르면 다음과 같은 창이 뜹니다.

저는 이 중에서 Tensorflow TFRecord 형식을 선택하였습니다. 하단의 show download code 옵션을 선택한 후 우측 하단의 continue 버튼을 누르면 코드가 생성될 텐데, 이를 복사하여 아래 코드의 녹색 부분에 붙여 넣으면 됩니다.

%mkdir /content/dataset
%cd /content/dataset
!curl -L "https://app.roboflow.com/your_dataset_link" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip

이제 코드를 실행하면 test/train/valid 각각에 대하여 tfrecord 파일과 label_map.pbtxt 파일이 생성되는 것을 확인할 수 있습니다.

학습을 위한 pretrained TF2 MobileNet v2 model 다운로드

SSD MobileNet v2를 학습하기 위해 pretrained TF2 MobileNet v2 모델을 feature extractor로 사용하였습니다. 이를 위해 다음과 같이 모델의 weights를 다운로드하였습니다.

%cd /content
!wget http://download.tensorflow.org/models/object_detection/classification/tf2/20200710/mobilenet_v2.tar.gz
!tar -xvf mobilenet_v2.tar.gz
!rm mobilenet_v2.tar.gz

!wget https://raw.githubusercontent.com/tensorflow/models/master/research/object_detection/configs/tf2/ssd_mobilenet_v2_320x320_coco17_tpu-8.config
!mv ssd_mobilenet_v2_320x320_coco17_tpu-8.config mobilenet_v2.config

Trainig parameter 설정

모델 학습에서 가장 중요한 것은 파라미터의 값이라고 생각합니다. 파라미터의 값을 어떻게 설정하느냐에 따라서 모델 학습의 성능이 달라질 수 있습니다. 하지만 파라미터 값은 딱 정해져있는 것이 아니라, 일종의 trial error이기 때문에 여러 번 시도해보면서 적당한 값을 찾아야 합니다. 저는 다음과 같이 설정하였습니다.

num_classes = 60 # dataset의 클래스 갯수
batch_size = 16
num_steps = 100000
num_eval_steps = 1000

train_record_path = '/content/dataset/train/your_dataset.tfrecord'
test_record_path = '/content/dataset/test/your_dataset.tfrecord'
model_dir = '/content/training/' #원하는 경로로 설정해주세요.
labelmap_path = '/content/dataset/train/your_dataset_label_map.pbtxt'

pipeline_config_path = 'mobilenet_v2.config'
fine_tune_checkpoint = '/content/mobilenet_v2/mobilenet_v2.ckpt-1'

이때 녹색으로 표시되어 있는 train_recode_path와 test_record_path, labelmap_path는 해당 파일이 위치해있는 경로로 수정해 줍니다. 그리고 만약 모델을 이어서 학습시키는 경우, fine_tune_checkpoint 경로를 이어서 학습시키고자 하는 체크포인트 파일로 수정해주면 됩니다.

config file 수정

이제 다음 코드를 수행하여 config file을 수정하도록 하겠습니다.

import re

%cd /content/
with open(pipeline_config_path) as f:
    config = f.read()

with open(pipeline_config_path, 'w') as f:

  # Set labelmap path
  config = re.sub('label_map_path: ".*?"', 
             'label_map_path: "{}"'.format(labelmap_path), config)
  
  # Set fine_tune_checkpoint path
  config = re.sub('fine_tune_checkpoint: ".*?"',
                  'fine_tune_checkpoint: "{}"'.format(fine_tune_checkpoint), config)
  
  # Set train tf-record file path
  config = re.sub('(input_path: ".*?)(PATH_TO_BE_CONFIGURED/train)(.*?")', 
                  'input_path: "{}"'.format(train_record_path), config)
  
  # Set test tf-record file path
  config = re.sub('(input_path: ".*?)(PATH_TO_BE_CONFIGURED/val)(.*?")', 
                  'input_path: "{}"'.format(test_record_path), config)
  
  # Set number of classes.
  config = re.sub('num_classes: [0-9]+',
                  'num_classes: {}'.format(num_classes), config)
  
  # Set batch size
  config = re.sub('batch_size: [0-9]+',
                  'batch_size: {}'.format(batch_size), config)
  
  # Set training steps
  config = re.sub('num_steps: [0-9]+',
                  'num_steps: {}'.format(num_steps), config)
  
  f.write(config)

딥러닝 모델 training

이제 실제로 모델을 학습하도록 하겠습니다. 아래 코드를 수행하면 설정한 파라미터에 따라 학습을 진행하게 됩니다.

!python /content/models/research/object_detection/model_main_tf2.py \
    --pipeline_config_path={pipeline_config_path} \
    --model_dir={model_dir} \
    --alsologtostderr \
    --num_train_steps={num_steps} \
    --sample_1_of_n_eval_examples=1 \
    --num_eval_steps={num_eval_steps}

학습이 진행되면서 각 step마다 loss 값이 출력되는 것을 확인할 수 있습니다. 그리고 model_dir로 지정해놓은 경로에 1000 step마다 체크포인트가 생성되는데, 만약 런타임이 끊기거나 학습이 갑작스레 중단된 경우 이 체크포인트를 이용하여 모델을 이어서 학습할 수 있습니다.

딥러닝 모델 evaluation

tensorflow 2의 경우에는 training을 하는 동시에 evaulation을 수행하는 기능을 지원하지 않습니다. 따라서 학습하는 동시에 evaluation을 수행하고 싶다면, evaluation을 위한 코드를 다른 터미널에서 병렬적으로 수행해주어야 합니다.

 

이를 위해 구글 코랩에서 새로운 노트를 만들어줍니다. 이 노트에도 앞서 설치하였던 각종 api와 dataset을 모두 import 해주고, training 코드 대신 evaluation을 위한 코드를 실행해줍니다.

!python /content/models/research/object_detection/model_main_tf2.py \
    --pipeline_config_path={pipeline_config_path} \
    --model_dir={model_dir} \
    --checkpoint_dir={model_dir}

이때, model_dir을 training 노트와 evlaluation 노트가 공유할 수 있도록 해주어야 합니다. 그래야 training을 수행하면서 생성되는 ckpt 파일을 evaluation에서 사용할 수 있으니까요.

 

이를 위해 저는 model_dir을 구글 드라이브의 경로로 설정해주었습니다. training을 하면서 구글 드라이브에 체크포인트 파일이 생성되면, 그 ckpt 파일을 이용해 다른 노트에서 동시에 evaluation을 수행할 수 있게 해 주기 위해서입니다.

 

이때, 한 가지 유의할 점이 있습니다. 구글 코랩 무료 사용 시, 두 개 이상의 노트에서 동시에 gpu를 사용할 수 없습니다. 그래서 저는 학습을 위한 노트에서는 런타임 유형을 gpu로 설정해주었고, eval을 위한 노트에서는 런타임 유형을 none으로 설정해주었습니다. evaluation에서는 딱히 gpu가 필요하지 않으니까요.

tensorboard 확인

학습이 완료되었다면, 이제 각종 성능 지표들을 그래프로 확인해보겠습니다. 이를 위해 아래 코드를 이용해서 텐서보드를 로드해줍니다.

%load_ext tensorboard
%tensorboard --logdir '/content/drive/MyDrive/model_dir_경로'

그러면 다음과 같이 텐서보드가 로드되고, train과 eval에 해당하는 각각의 loss 그래프를 한눈에 확인할 수 있습니다.

저 같은 경우에는 이미 어느정도 학습된 모델을 이어서 학습시켰기 때문인지 드라마틱한 감소폭을 보이지는 않고, 중간중간 값이 튀기도 하면서 변동적인 모습을 보이네요. 그래도 전체적으로 보면 train loss와 validation loss가 함께 잘 감소하고 있는 것을 확인할 수 있습니다.

 

텐서보드를 통해 loss 이외에도 accuracy, recall 값 등도 확인할 수 있습니다.

Inference Graph export

이제 trained model을 inference graph로 export 및 다운로드해 줍니다.

output_directory = 'inference_graph'

!python /content/models/research/object_detection/exporter_main_v2.py \
    --trained_checkpoint_dir {model_dir} \
    --output_directory {output_directory} \
    --pipeline_config_path {pipeline_config_path}

!zip -r /content/saved_model.zip /content/inference_graph/saved_model/

from google.colab import files
files.download("/content/saved_model.zip")

그러면 다음과 같은 구성을 가진 saved_model 파일을 얻을 수 있습니다.

inference_graph - saved model --- assets --- variables ----- variables.data-00000-of-00001 ----- variables.index --- saved_model.pb

모델을 이용하여 Inference 수행

그럼 이제 학습한 모델이 이미지에서 원하는 객체를 잘 인식하는지 테스트해보도록 하겠습니다. 우선 모델을 불러오도록 하겠습니다.

!wget https://raw.githubusercontent.com/hugozanini/object-detection/master/inferenceutils.py
from inferenceutils import *

output_directory = 'inference_graph/'

category_index = label_map_util.create_category_index_from_labelmap(labelmap_path, use_display_name=True)
tf.keras.backend.clear_session()
model = tf.saved_model.load(f'/content/{output_directory}/saved_model')

테스트할 이미지를 불러오도록 하겠습니다. roboflow에서 이번에는 csv format으로 데이터셋을 불러옵니다.

%mkdir /content/csv_dataset
%cd /content/csv_dataset
!curl -L "https://app.roboflow.com/ds/wVqerq26Oc?key=7rSF6O4vnq" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip

아래 코드를 실행하여 원하는 이미지에 대해 inference를 수행할 수 있습니다.

images = ['/content/csv_dataset/test/your_image1.jpg',
          '/content/csv_dataset/test/your_image2.jpg',
          '/content/csv_dataset/test/your_image3.jpg']

for image_name in images:
  
  image_np = load_image_into_numpy_array(image_name)
  output_dict = run_inference_for_single_image(model, image_np)
  vis_util.visualize_boxes_and_labels_on_image_array(
      image_np,
      output_dict['detection_boxes'],
      output_dict['detection_classes'],
      output_dict['detection_scores'],
      category_index,
      instance_masks=output_dict.get('detection_masks_reframed', None),
      use_normalized_coordinates=True,
      line_thickness=8)
  display(Image.fromarray(image_np))

다음과 같이 객체 인식 결과 이미지를 확인할 수 있습니다.

object detected image 1

object detected image 2

object detected image 3

학습한 음식에 대해서는 곧잘 인식을 수행하는 것을 확인할 수 있습니다. 다만 첫 번째 사진의 경우, 객체를 정확히는 인식하지 못하는 것을 보아 추가적인 학습이 필요해 보입니다.

---

3. tensorflow.js 모델로 변환

학습한 모델을 크롬 익스텐션에서 사용하기 위해서는 자바 스크립트 기반의 딥러닝 모델인 tfjs로의 변환이 필요합니다. 이를 위해 위에서 얻은 saved_model을 이용하여 tensorflow.js converter를 통해 변환을 수행하도록 하겠습니다.

우선 터미널에 다음과 같은 코드를 입력해 tensorflow.js converter를 설치합니다.

 pip install tensorflowjs[wizard]

설치를 완료하였으며 이제 다음 코드를 통해 converter를 실행합니다.

tensorflowjs_wizard

여러 옵션을 직접 선택하여 원하는 대로 모델을 convert 할 수 있습니다. 저는 아래와 같은 옵션으로 모델을 변환하였습니다.

이러한 과정을 통해 tfjs 모델을 얻을 수 있습니다. web_model 파일에 model.json과 sharded weigths file이 들어있는 것을 확인할 수 있습니다.

이렇게 얻은 tensorflow.js 모델을 이용하여 웹 브라우저 환경에서 객체 인식을 수행할 수 있습니다.

---

이번 게시글에서는 tensorflow를 이용하여 딥러닝 모델을 학습시키고, 웹 브라우저 적용을 위해 tensorflow.js로 변환하는 방법을 다루었습니다.

 

기회가 된다면 다음 게시글에서는 모델을 크롬 익스텐션에 적용하여 유튜브 영상에 대한 Object Detection 방법을 다루도록 하겠습니다.