Python PyTorch SyncBatchNorm用法及代碼示例

本文簡要介紹python語言中 torch.nn.SyncBatchNorm 的用法。

用法:

class torch.nn.SyncBatchNorm(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True, process_group=None, device=None, dtype=None)

參數：

• num_features-C 來自大小為 (N, C, +) 的預期輸入

• eps-加到分母上的值，以保證數值穩定性。默認值：1e-5

• momentum-用於 running_mean 和 running_var 計算的值。可以設置為None 用於累積移動平均(即簡單平均)。默認值：0.1

• affine-一個布爾值，當設置為 True 時，此模塊具有可學習的仿射參數。默認值：True

• track_running_stats-一個布爾值，當設置為 True 時，此模塊跟蹤運行均值和方差，當設置為 False 時，此模塊不跟蹤此類統計信息，並將統計緩衝區 running_mean 和 running_var 初始化為 None .當這些緩衝區為 None 時，此模塊始終使用批處理統計信息。在訓練和評估模式下。默認值：True

• process_group-統計信息的同步發生在每個進程組中。默認行為是全局同步

如本文所述，在 N 維輸入(具有額外通道維度的 [N-2]D 輸入的小批量)上應用批量歸一化批量標準化：通過減少內部協變量偏移來加速深度網絡訓練.

y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta

均值和標準差是針對同一過程組的所有小批量按維度計算的。 \gamma 和 \beta 是大小為 C 的可學習參數向量(其中 C 是輸入大小)。默認情況下，\gamma 的元素是從 \mathcal{U}(0, 1) 中采樣的，\beta 的元素設置為 0。標準差是通過有偏估計器計算的，相當於 torch.var(input, unbiased=False) 。

同樣默認情況下，在訓練期間，該層會繼續對其計算的均值和方差進行估計，然後在評估期間將其用於歸一化。運行估計保持默認 momentum 為 0.1。

如果 track_running_stats 設置為 False ，則該層不會繼續運行估計，並且在評估期間也會使用批處理統計信息。

注意

這個momentum 參數不同於優化器類中使用的參數和傳統的動量概念。在數學上，這裏運行統計的更新規則是 \hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t ，其中 \hat{x} 是估計的統計量，而 x_t 是新的觀察值。

因為批量歸一化是針對C 維度中的每個通道進行的，計算(N, +) 切片的統計信息，所以通常將其稱為體積批量歸一化或Spatio-temporal 批量歸一化。

目前SyncBatchNorm 僅支持每個進程單個 GPU 的 DistributedDataParallel (DDP)。在使用 DDP 包裝網絡之前，使用 torch.nn.SyncBatchNorm.convert_sync_batchnorm() 將 BatchNorm*D 層轉換為 SyncBatchNorm。

形狀：

• 輸入：(N, C, +)

• 輸出：(N, C, +)(與輸入的形狀相同)

注意

batchnorm 統計信息的同步僅在訓練時發生，即當設置 model.eval() 或 self.training 為 False 時禁用同步。

例子：

>>> # With Learnable Parameters
>>> m = nn.SyncBatchNorm(100)
>>> # creating process group (optional)
>>> # ranks is a list of int identifying rank ids.
>>> ranks = list(range(8))
>>> r1, r2 = ranks[:4], ranks[4:]
>>> # Note: every rank calls into new_group for every
>>> # process group created, even if that rank is not
>>> # part of the group.
>>> process_groups = [torch.distributed.new_group(pids) for pids in [r1, r2]]
>>> process_group = process_groups[0 if dist.get_rank() <= 3 else 1]
>>> # Without Learnable Parameters
>>> m = nn.BatchNorm3d(100, affine=False, process_group=process_group)
>>> input = torch.randn(20, 100, 35, 45, 10)
>>> output = m(input)

>>> # network is nn.BatchNorm layer
>>> sync_bn_network = nn.SyncBatchNorm.convert_sync_batchnorm(network, process_group)
>>> # only single gpu per process is currently supported
>>> ddp_sync_bn_network = torch.nn.parallel.DistributedDataParallel(
>>>                         sync_bn_network,
>>>                         device_ids=[args.local_rank],
>>>                         output_device=args.local_rank)