메시-텐서플로우를 이용한 초 고해상도 이미지 분석
2020년 2월 28일 (금) 구글 AI 리서치 분석
심층 신경망 모델은 대부분의 최신 이미지 분석 및 자연어 처리 알고리즘의 백본을 형성합니다. 최근에 데이터 및 모델의 병렬 처리와 같은 대규모 딥러닝 기술이 개발됨에 따라 대형 CNN (Convolutional Neural Network) 모델을 몇 분 만에 수백만 개의 이미지 데이터에 대해 학습할 수 있습니다. 그러나 최대 108 픽셀을 가질 수 있는 3D CT (컴퓨터 단층 촬영, computed tomography) 이미지와 같은 초 고해상도 이미지에 CNN 모델을 적용하는 것은 여전히 어려운 일입니다. 기존 기술을 사용하면 프로세서는 여전히 최소 32GB의 부분 중간 데이터를 호스팅해야 하지만 개별 GPU 또는 TPU에는 일반적으로 12-32GB 메모리만 있습니다. 일반적인 솔루션은 이미지 패치를 서로 별도로 처리하여 정보 손실로 인해 복잡한 구현 및 최적화 성능이 이하로 제공됩니다.
Mayo Clinic과 공동 작업으로 “공간 분할을 사용한 고해상도 의료 이미지 분석(High Resolution Medical Image Analysis with Spatial Partitioning” 논문에서, Mesh-TensorFlow 프레임워크를 사용하여 방대한 데이터 및 모델 병렬 처리의 경계를 넓히고 실제 실행 가능성을 위해 입력 해상도를 손상시키지 않고 이 기술을 초 고해상도 이미지 분석에 어떻게 사용할 수 있는지 보여줍니다.
인접 파티션 간의 관계를 유지하기 위해 공간 파티션에서 컨볼루션 연산을 처리하기 위한 후광 변환(halo exchange) 알고리즘을 구현합니다. 결과적으로, 256-way 모델 병렬 처리로 초 고해상도 이미지 (각 차원에서 512 픽셀의 3D 이미지)에서 3D U-Net을 학습할 수 있습니다. 또한 광범위한 리서치 커뮤니티에서 사용할 수 있도록 GPU와 TPU를 위한 Mesh-TensorFlow 기반 프레임 워크를 추가로 오픈소스화 했습니다.
Mesh-TensorFlow를 사용한 데이터 및 모델 병렬 처리
우리의 구현은 쉽고 효율적인 데이터 및 모델 병렬 처리를 위해 Mesh-TensorFlow 프레임워크를 기반으로 하며, 사용자 정의 이미지 레이아웃에 따라 사용자가 여러 장치의 메쉬에 텐서들을 분할할 수 있습니다. 예를 들어, 사용자는 계산 장치의 메시를 프로세서 당 2개의 코어와 함께 총 256개의 프로세서에 대해 16행 x 16열로 제공할 수 있습니다. 그런 다음 이미지의 공간 차원 x를 프로세서 행에 매핑하고 공간 차원 y를 프로세서 열에 매핑하고 배치 차원(즉, 동시에 처리할 이미지 세그먼트 수)을 코어에 매핑하는 레이아웃을 정의합니다. 훈련 배치(training batch)의 분할 및 분산은 사용자가 구현 세부 사항에 대해 걱정할 필요 없이 텐서 레벨에서 Mesh-TensorFlow에 의해 구현됩니다. 아래 그림은 간단한 예를 가진 개념을 보여줍니다.
[그림 1] 초 고해상도 이미지의 공간 분할 (이 경우 3D CT 스캔)Halo Exchange를 사용한 공간 분할
이미지에서 실행되는 컨볼루션 작용(convolution operation)은 종종 프레임 에지를 넘어 확장되는 필터를 적용합니다. 단일 이미지를 다룰 때 이 문제를 해결할 수 있는 방법이 있지만 표준 방법에서는 프레임 에지를 넘어 세그먼트화 된 이미지 정보가 여전히 관련될 수 있다는 점을 고려하지 않습니다. 정확한 결과를 얻으려면 공간적으로 분할되고 프로세서 간에 재배포된 이미지에 대한 컨볼루션 작업은 각 이미지 세그먼트의 이웃을 고려해야 합니다.
하나의 가능한 해결책은 각 공간 파티션에 겹치는 영역을 포함하는 것입니다. 그러나 많은 후속 컨볼루션 레이어가 있을 가능성이 높고 각 레이어에 겹침이 발생하므로 겹침이 상대적으로 커집니다. 실제로 대부분의 경우 겹침이 전체 이미지를 덮을 수 있습니다. 또한 모든 중첩 영역은 시작부터 첫 번째 계층에 포함되어야 하며, 이는 해결하려는 메모리 제약 조건으로 이어질 수 있습니다.
우리의 솔루션은 완전히 다릅니다. 우리는 후광 교환이라는 데이터 통신 단계를 구현했습니다. 모든 컨볼루션 작업 전에 각 공간 파티션은 주변과 마진을 교환(수신 및 전송)하여 마진에서 이미지 세그먼트를 효과적으로 확장합니다. 컨볼루션 작업은 각 장치에 로컬로 적용됩니다. 이것은 전체 이미지에 대한 컨볼루션 결과가 공간 분할과 동일하거나 공간 분할 없이 유지되도록 합니다.
[그림 2] Halo 교환은 파티션 간 컨볼루션의 이미지 세그먼트 가장자리를 올바르게 처리하도록 합니다.Proof of Concept-간 종양 CT 스캔의 분할
그런 다음이 프레임워크를 간 종양(LiTS 벤치 마크)의 3D CT 스캔을 분할하는 작업에 적용했습니다. 평가 척도의 경우, Sørensen-Dice 계수(coefficient)를 사용합니다. 이 계수는 0.0에서 1.0 사이이며, 점수가 0인 세그먼트와 실제(ground true) 종양 영역 간에 겹치지 않음을 나타내고 1은 완벽한 일치를 나타냅니다. 아래에 표시된 결과는 높은 데이터 해상도로 더 나은 결과를 얻을 수 있음을 보여줍니다. 전체 5123 해상도 (각 x, y, z 방향에서 512 픽셀)를 사용할 때 수익이 감소하는 경향이 있지만, 이 작업은 초 고해상도 이미지 분석의 가능성을 열어줍니다.
[그림 3] 더 높은 해상도 데이터는 더 나은 분할 정확도를 제공합니다.결론
기존 데이터 및 모델 병렬 처리 기술을 사용하면 수십억 개의 매개 변수로 신경망을 학습할 수 있지만, ~108 픽셀 이상의 입력 이미지는 처리할 수 없습니다. 이 연구에서는 이러한 초 고해상도 이미지에 대한 CNN의 적용 가능성을 살펴보고 유망한 결과를 보여줍니다. 우리의 Mesh-TensorFlow 기반 구현은 GPU와 TPU 모두에서 작동하며 릴리스 된 코드를 사용하여 이전에는 불가능했던 일부 작업에 가능한 솔루션을 제공하기를 희망합니다.
감사의 말
초기 3D U-net 모델 및 학습 데이터를 제공한 파트너인 Mayo Clinic의 Panagiotis Korfiatis 박사 및 Daniel Blezek 박사에게 감사드립니다. Mayo Clinic과의 POC 작업에 대해 Greg Mikels에게 감사합니다. 논문의 공동 저자, 특히 Noam Shazeer에게 특별한 감사를 드립니다.
원본 제목: Mesh-TensorFlow를 이용한 초 고해상도 이미지 분석(Ultra-High Resolution Image Analysis with Mesh-TensorFlow)
게시자 : Le Hou 및 Youlong Cheng, Google Research 소프트웨어 엔지니어
원본 링크: https://ai.googleblog.com/2020/02/ultra-high-resolution-image-analysis.html
* Mesh TensorFlow 오픈소스: https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/experimental/model_executor.py
* Mesh 오픈소스: https://github.com/tensorflow/mesh/tree/master/mesh_tensorflow/experimental
이 블로그는 2020년 2월 28일(금), Google AI Research Blog 기사를 영어에서 우리나라 말로 번역한 것입니다. 또한 이 번역 글은 정보 공유 목적으로만 작성했으므로 어떠한 상업용으로 사용할 수 없으며, 원본 저작물 모두 구글에게 저작권이 있음을 알려 드립니다. (First Draft Version)
