Cloud TPU로 큰 이미지 처리

이 문서에서는 Cloud TPU 및 TPUEstimator를 사용한 공간 분할로 매우 큰 이미지와 동영상의 학습을 지원하는 방법을 설명합니다.

Cloud TPU 공간 분할 기능을 사용하면 여러 TPU 칩에서 단일 이미지를 분할할 수 있습니다. 공간 분할을 사용하면 액셀러레이터 칩의 단일 코어에 사용 가능한 메모리에 맞지 않을 정도로 큰 입력 이미지를 사용하여 실행할 수 있도록 모델을 쉽게 확장할 수 있습니다.

이 기능은 활성화를 여러 TPU 코어에 분산하여 2개, 4개, 8개 또는 16개의 코어를 학습에 사용하도록 모델을 확장할 수 있습니다. 학습에 사용할 수 있는 프로세서의 수를 확장하면 더 높은 정확도와 더 효율적인 학습 성능을 제공하는 다운샘플링 이미지 데이터가 필요하지 않습니다.

다음과 같은 태스크에 공간 분할을 적용할 수 있습니다.

  • 3D 컴퓨터 단층 촬영(CT) 스캔 이미지 세분화
  • 동영상 콘텐츠 분석
  • 이미지 데이터를 다운샘플링하지 않고 자율 주행을 위해 객체 감지

TPUEstimator를 사용한 공간 분할 사용 설정

TensorFlow에서 TensorFlow용 XLA 최적화 컴파일러는 모든 Cloud TPU 코어 간의 통신을 자동으로 처리합니다. 모델의 공간 분할을 사용 설정하기 위해 코드를 변경할 필요가 없습니다. TPUEstimator가 공간 분할 API를 지원하므로, TPUConfig에서 각 입력 텐서를 분할하는 방법을 구성하기만 하면 됩니다.

이미지

예

다음 코드는 이미지 분류 모델에서 4방향 공간 분할의 TPUConfig 예시를 보여줍니다. 텐서는 높이 차원을 따라 네 부분으로 나뉩니다(텐서에 모양[배치, 높이, 너비, 채널]이 있다고 가정).

tpu_config=tpu_config.TPUConfig(
    iterations_per_loop=100,
    num_cores_per_replica=4,
    per_host_input_for_training=tpu_config.InputPipelineConfig.PER_HOST_V2,
    input_partition_dims=[[1, 4, 1, 1], None]]

공간 분할의 경우 입력 파이프라인은 tf.data 형식이어야 하고 per_host_input_for_training 값(train_batch_size)은 PER_HOST_V2로 설정되어야 합니다. 지정한 num_cores_per_replica는 이미지를 분할할 수있는 파티션의 최대 개수를 결정합니다.

input_partition_dims 값은 입력 텐서를 분할하는 방법을 설명하는 feature_partition_dims와 label_partition_dims라는 두 개의 요소가 있는 목록을 제공합니다. feature_partition_dims 및 label_partition_dims의 구조는 input_fn의 기능 및 라벨의 구조와 일치해야 합니다. 라벨이 분할되지 않도록 라벨 파티션에 '없음'을 지정하세요.

공간 분할을 사용하여 참조 모델 실행

2D 객체 감지

RetinaNet은 경계 상자를 사용하여 이미지에서 객체를 찾고 식별된 객체를 분류하는 객체 감지 모델입니다. 하나의 Cloud TPU 코어(기기당 8개의 배치)에 맞는 가장 큰 이미지는 1280x1280이므로 매우 큰 이미지를 사용하여 모델을 학습시키기는 어려운 과제입니다. 그림과 같이 8개의 TPU 코어에 모델을 공간적으로 분할하여 4배 더 큰 이미지를 학습시킬 수 있습니다.

이미지 크기 TPU 유형 TPU 코어 전역 배치 크기 단계 시간
1280x1280 v3-8 8 64 910ms
2560x2560 v3-64 64 64 822ms

3D 이미지 세분화

3D UNet은 의학 이미지 분야에서 널리 사용되는 고밀도 3D 세분화 모델입니다. CT 이미지의 원래 해상도는 단일 TPU 코어에 맞지 않는 256x256x256 크기이므로 연구자가 이미지를 다운샘플링해야 하는 경우가 많습니다. TPU 공간 분할을 사용하면 16방향 공간 분할을 통해 원본 해상도 이미지에 직접 맞출 수 있습니다.

이미지 크기 TPU 유형 TPU 코어 전역 배치 크기 단계 시간
128x128x128 v3-8 8 32 3.428s
2560x2560 v3-64 64 32 3.02s

추가 리소스