用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部署。