用NVIDIA Tensor Cores和TensorFlow 2加速醫學圖像分割

用NVIDIA Tensor Cores和TensorFlow 2加速醫學圖像分割

Accelerating Medical Image Segmentation with NVIDIA Tensor Cores and TensorFlow 2

醫學圖像分割是當前學術界研究的熱點。這方面正在進行的挑戰、競賽和研究項目的數量證實了這一點，這些項目的數量只是逐年上升。在解決這一問題的各類方法中，U-Net已經成爲許多2D和3D分割任務的最佳解決方案的骨幹。這是由於簡單性、多功能性和有效性。             

當實踐者面臨一個新的分割任務時，第一步一般是使用現有的U-Net實現做爲主幹。可是，隨着TensorFlow 2.0的到來，缺乏現成的解決方案。如何有效地將模型轉換到TensorFlow 2.0以利用新功能，同時仍然保持頂級硬件性能並確保最早進的精度？

 

Figure 1. Example of a serial section Transmission Electron Microscopy image (ssTEM) and its corresponding segmentation

U-Net for medical image segmentation

U-Net首先由Olaf Ronneberger、Philip Fischer和Thomas Brox引入。在本文中，U-Net：用於生物醫學圖像分割的卷積網絡。U-Net容許以高精度和高性能無縫分割2D圖像。能夠用來解決許多不一樣的分割問題。             
圖2顯示了U-Net模型及其不一樣組件的構造。U-Net由一條壓縮和一條擴展路徑組成，其目的是經過卷積和池操做的結合，在其最核心的部分創建一個瓶頸。在這個瓶頸以後，經過卷積和上採樣相結合的方法重建圖像。跳躍鏈接的目的是幫助梯度的反向流動，以改善訓練。

 

Figure 2. The architecture of a U-Net model. Source: U-Net: Convolutional Networks for Biomedical Image Segmentation

U-Net擅長的任務一般被稱爲語義分割，須要用對應的類來標記圖像中的每一個像素，以反映所表示的內容。由於對圖像中的每一個像素都執行此操做，因此此任務一般稱爲密集預測。             

在語義分割的狀況下，預測的預期結果是高分辨率圖像，一般與被饋送到網絡的圖像具備相同的維度，其中每一個像素被標記到相應的類。廣義地說，語義分割只是一種像素級、多類分類的形式。             

儘管U-Net主要用於語義分割，但基於U-Net的網絡在目標檢測或實例分割等任務中取得很好的效果並很多見。

Migrating to TensorFlow 2.0 with performance in mind

一般，使用深度學習模型的第一步是創建一個讓感到溫馨的基線。在NVIDIA深度學習示例Github存儲庫中，能夠找到最流行的深度學習模型的實現。這些實現幾乎涵蓋了每一個領域和框架，並提供了普遍的基準，以確保最佳的準確性和性能。所以，不管是從業者仍是研究者，都是最佳起點。             

在這些實現中，能夠找到U-Net，能夠在TensorFlow 1.x和TensorFlow 2.0中找到。可是，遷移到TensorFlow的最新版本須要遵循哪些步驟？

A new way to run models

在這個新版本的TensorFlow中，最顯著的變化之一是在使用會話到函數調用之間的切換。到目前爲止，將指定要調用的輸入和函數，並指望返回模型的輸出。而後在會話.run調用，以下代碼示例所示：

TensorFlow 1.X

outputs = session.run(f(placeholder), feed_dict={placeholder: input})

TensorFlow 2.0

outputs = f(input)

這是可能的，由於TensorFlow 2.0中默認啓用了緊急執行。這改變了與TensorFlow交互的方式，由於緊急執行是一個當即評估操做的命令式編程環境。提供了一些好處，例如更直觀的界面、更容易調試、天然的控制流，但代價是性能較差。一般被推薦用於研究和實驗。             

AutoGraph             

要在使用TensorFlow 2.0的模型中實現生產級性能，必須使用AutoGraph（AG）。這個Tensorflow 2.0特性容許經過使用decorator使用天然的Python語法編寫Tensorflow圖形代碼@tf.函數，以下代碼示例所示：

@tf.function

def train_step(features, targets, optimizer):

With tf.GradientTape() as tape:

        predictions = model(features)

        loss = loss_fn(predictions, targets)

    vars = model.trainable_variables

gradients = tape.gradient(loss, vars)

optimizer = apply_gradients(zip(gradients, vars))

儘管AG仍有侷限性，但帶來的性能改進是顯而易見的。有關如何使用tf.函數和AG，參見TensorFlow 2.0指南。

Maximizing Tensor Core usage

混合精度是不一樣數值精度在計算方法中的綜合運用。混合精度訓練經過以半精度格式執行操做，同時以單精度存儲最小信息以在網絡的關鍵部分保留儘量多的信息，從而顯著加快計算速度。             

在Volta和Turing體系結構中引入張量核以後，能夠經過切換到混合精度來體驗顯著的訓練加速：在最嚴格的數學模型體系結構上，整體加速高達3倍。使用混合精度訓練須要兩個步驟：

1.       Porting the model to use the FP16 data type where appropriate.

2.       Adding loss scaling to preserve small gradient values.

For more information, see the following resources:

• Mixed Precision Training paper 
• Training With Mixed Precision documentation.
• Mixed-Precision Training of Deep Neural Networks post.
• Using TF-AMP from the TensorFlow User Guide.

Enabling automatic mixed precision training

要在TensorFlow 2.0中啓用自動混合精度（AMP），必須對代碼應用如下更改。             

設置Keras混合精度策略：                   

tf.keras.mixed_precision.experimental.set_policy('mixed_float16')

在優化器上使用損失縮放包裝器。默認狀況下，能夠選擇動態損耗縮放：

optimizer = tf.keras.mixed_precision.experimental.LossScaleOptimizer(optimizer, "dynamic")

確保使用縮放損失計算梯度：

loss = loss_fn(predictions, targets)scaled_loss = optimizer.get_scaled_loss(loss)

vars = model.trainable_variables

scaled_gradients = tape.gradient(scaled_loss, vars)

gradients = optimizer.get_unscaled_gradients(scaled_gradients)

optimizer.apply_gradients(zip(gradients, vars))

Enabling Accelerated Linear Algebra

加速線性代數（XLA）是一種特定於領域的線性代數編譯器，能夠加速TensorFlow模型，而沒必要更改源代碼。其結果是速度和內存使用的提升：在啓用XLA以後，大多數內部基準測試的運行速度提升了1.1-1.5倍。             

要啓用XLA，請在優化器中設置實時（JIT）圖形編譯。能夠經過對代碼進行如下更改來完成此操做：

tf.config.optimizer.set_jit(True)

U-Net in TensorFlow 2.0

在NVIDIA深度學習示例GitHub存儲庫中，能夠找到使用TensorFlow 2.0的U-Net實現。這個實現包含了全部必要的部分，不只能夠將U-Net移植到新版本的Google框架中，還可使用

tf.estimator.Estimator.

除了前面描述的與模型性能相關的必要更改外，請遵循如下步驟以確保模型徹底符合新API：

• API cleanup
• Data loading
• Model definition

API cleanup

因爲TensorFlow 2.0中的許多API操做已被刪除或已更改位置，第一步是使用v2升級腳本將不推薦的調用替換爲新的等效調用：

tf_upgrade_v2 \  

  --intree unet_tf1/ \

  --outtree unet_tf2/ \

  --reportfile report.txt

儘管此轉換腳本自動化了大部分過程，但根據報告文件中捕獲的建議，仍須要手動進行一些更改。有關更多信息，TensorFlow提供瞭如下指南，自動將代碼升級到TensorFlow 2。             

數據加載             

用於訓練模型的數據管道與用於TensorFlow 1.x實現的數據管道相同。是使用tf.data.Dataset數據集API操做。數據管道加載圖像並使用不一樣的數據加強技術對其進行轉換。有關詳細信息，請參見data_loader.py腳本。             

模型定義             

這個新版本的TensorFlow鼓勵將代碼重構成更小的函數並模塊化不一樣的組件。其中之一是模型定義，如今能夠經過子類化來執行tf.keras.Model型:

class Unet(tf.keras.Model):    """ U-Net: Convolutional Networks for Biomedical Image Segmentation

    Source:

        https://arxiv.org/pdf/1505.04597

    """

    def __init__(self):

        super().__init__(self)

        self.input_block = InputBlock(filters=64)

        self.bottleneck = BottleneckBlock(1024)

        self.output_block = OutputBlock(filters=64, n_classes=2)

        self.down_blocks = [DownsampleBlock(filters, idx)

                            for idx, filters in enumerate([128, 256, 512])]

        self.up_blocks = [UpsampleBlock(filters, idx)

                          for idx, filters in enumerate([512, 256, 128])]

    def call(self, x, training=True):

        skip_connections = []

        out, residual = self.input_block(x)

        skip_connections.append(residual)

        for down_block in self.down_blocks:

            out, residual = down_block(out)

            skip_connections.append(residual)

        out = self.bottleneck(out, training)

        for up_block in self.up_blocks:

            out = up_block(out, skip_connections.pop())

        out = self.output_block(out, skip_connections.pop())

        return tf.keras.activations.softmax(out, axis=-1)

Performance results

根據上述更改，運行模型以驗證TensorFlow2.0上Tensor Cores提供的加速。有三種不一樣的特性影響性能：

• The use of TensorFlow AG
• The use of Accelerated Linear Algebra (XLA)
• The use of automatic mixed precision (AMP) using Tensor Cores

Training performance

結果是經過ensorflow:20.02-tf-py3下一代NVIDIA DGX-1上帶有（8x V100 16G）GPU的容器。性能數字（以每秒項/圖像爲單位）在1000次迭代中平均，不包括前200個預熱步驟。

 

Table 1. Single- and multi-GPU training performance for FP32 and mixed precision. The speedup is the ratio of images per second processed in mixed precision compared to FP32.

在TensorFlow 2.0中，對於單GPU訓練，使用混合精度的例子模型可以達到每秒圖像測量的2.89x加速比；對於多GPU訓練，使用混合精度的例子模型可以達到每秒圖像測量的2.7x加速比，使用混合精度的例子模型可以達到幾乎完美的弱比例因子。有關更多信息，請參閱遵循的步驟。             

對TensorFlow2.0中可用的不一樣特性如何影響訓練階段性能的額外瞭解。提升模型吞吐量的最有效方法是啓用AMP，包括此功能在內的全部設置都是性能最好的。可是，當使用混合精度進行訓練時，XLA只有在使用AG時才能提供提高。AMP、XLA和AG的組合提供了最佳的結果。

Inference performance

經過運行tensorflow:20.02-tf-py3下一代NVIDIA DGX-1上的容器，帶有（1x V100 16G）GPU。吞吐量以每秒圖像幀爲單位報告。每批報告的延遲時間以毫秒爲單位。

 

表2. FP32和混合精度的單GPU推理性能。加速比是每秒以混合精度處理的圖像與FP32的比率。

該模型在TensorFlow 2.0推理機上使用混合精度，可達到3.19x的加速比。有關更多信息，請參閱推理性能基準測試步驟。

TensorFlow2.0中可用的不一樣特性如何影響推理階段性能的額外細節。最重要的性能提高出如今啓用放大器時，由於張量核大大加速了預測。由於模型已通過訓練，因此能夠經過啓用AG來禁用急切執行，這將提供額外的加強。

Next steps

可使用NVIDIA gpu開始利用TensorFlow 2.0中的新功能。TensorFlow2.0極大地簡化了模型的訓練過程，而且在最新版本的框架中比以往任什麼時候候都更容易利用Tensor核心。             

在Deep Learning Examples存儲庫中，將找到25個以上最流行的Deep Learning模型的開源實現。有TensorFlow、Pythorch和MXNet。在這些例子中，是關於如何以最早進的精度和破紀錄的性能運行這些模型的逐步指南。             

全部實現都附帶了有關如何執行混合精度訓練以使用張量核加速模型的說明。還以檢查點的形式分發這些模型。這些模型是徹底維護的，包括最新的功能，如多GPU訓練、使用DALI加載數據和TensorRT部署。