tensorflow 2.6, 2.7 에서는 tf.keras.layers.preprocessing 하위에 레이어들이 추가되었다. 이것을 이용해서 복잡한 tensorflow native api 를 사용하지 않더라도 간단한 몇몇 augmentation 을 빠르게 레이어처럼 끼워넣을 수 있게 되었다.
하지만 이 augmentation layer 들은 1개의 input 만 받을 수 있는데, 이렇게 되면 image 와 mask 둘 모두를 사용하는 나의 파이프라인에는 맞지 않게 된다. 이런 점에서 Albumentation 과 크게 대조적이다.
Deep learning based computer vision 에는 많은 태스크들이 있고, 각 태스크마다 augmentation 의 전략은 모두 다르기 때문에, 앞으로도 tensorflow 에서 segmentation task 를 위한 파이프라인을 따로 만들어주지는 않을 것이다. 그러니까 내가 직접 간단히 만들어 써야 한다.
일부 augmentation 은 사전에 적용된 데이터셋을 사용하면서 (참고1), 일부 augmentation 은 런타임에 돌아가도록 만든다. 아래는 구상도이다.
나는 segmentation model 이라는 segmentation 의 오픈소스 high level api 를 사용하고 있다.
우리의 segmentation model 은 model.fit(tf_dataset) 로 훈련한다. 이 model.fit(tf_dataset) 을 그대로 두고 사용하기 위해서는 도대체 어떻게 해야할까? tf_datset 은 (image, mask) tuple 로 이루어져 있다. model.fit(x=image, y=mask) 로 언패킹되어 들어가는 셈이다.
즉, 전체 파이프라인이 model.fit() 으로 실행되기 때문에, API 호출 형태를 변형시키지 않고 커스텀 파이프라인을 작동시키기 위해서는 model.fit() 을 손질해 주어야 한다.
model.fit 에는 model(data) 의 형태로 들어가기 때문에, model.call() 부분을 TPU 로 돌렸을 때 아무런 문제가 없이, 가속화가 잘 되는지 미리 확인해 보았고 별다른 문제가 없음을 확인했다 (참고2).
나의 전략은, 어차피 학습 단계에서 일어나는 일들은 train_step() 이라는 함수가 관장하고 (참고3), model.fit() 의 argument 들을 관리할 수 있기 때문에 이 함수를 수정하고자 했다. COLAB notebook 참고 (참고4).
to
참고
1.
2.