TensorFlow를 이용한 YOLO v1 논문 구현 #6 - model.py

이전 글: https://it-the-hunter.tistory.com/33

TensorFlow를 이용한 YOLO v1 논문 구현 #5 - utils.py

model.py

목표: keras subclassing 형태의 모델 구조 class 정의

---

필요한 모듈 import

import tensorflow as tf

---

상세 코드

class YOLOv1(tf.keras.Model):

keras model을 상속받는 class YOLOv1 정의

  def __init__(self, input_height, input_width, cell_size, boxes_per_cell, num_classes):
    super(YOLOv1, self).__init__()
    base_model = tf.keras.applications.InceptionV3(include_top=False, weights='imagenet', input_shape=(input_height, input_width, 3))
    base_model.trainable = True
    x = base_model.output

    #Global Average Pooling
    x = tf.keras.layers.GlobalAveragePooling2D()(x)
    output = tf.keras.layers.Dense(cell_size * cell_size * (num_classes + (boxes_per_cell*5)), activation=None)(x)
    model = tf.keras.Model(inputs=base_model.input, outputs=output)
    self.model = model
    #model structure 출력
    self.model.summary()

keras 모델 중 pre-train된 모델 InceptionV3를 가져오는 코드

하지만 base_model.trainable = True 코드를 넣어 레이어도 우리가 원하는 대로 train이 가능하게 함

---

  def call(self, x):
    return self.model(x)

input으로 넘어온 이미지에 대해 모델 구조에 대한 예측값 반환