SyncBatchNorm#

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

對 N 維輸入應用批歸一化。

N 維輸入是 [N-2] 維輸入的 mini-batch（具有額外的通道維度），如論文 Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift 中所述。

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

均值和標準差是按維度在所有同屬一個程序組的 mini-batch 中計算的。 $\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$ 是新觀察到的值。

由於 Batch Normalization 是針對 C 維度中的每個通道進行的，因此在 (N, +) 切片上計算統計量，因此通常稱之為 Volumetric Batch Normalization 或 Spatio-temporal Batch Normalization。

目前 SyncBatchNorm 只支援單 GPU/程序的 DistributedDataParallel (DDP)。在將網路包裝到 DDP 之前，使用 torch.nn.SyncBatchNorm.convert_sync_batchnorm() 將 BatchNorm*D 層轉換為 SyncBatchNorm。

引數

num_features (int) – $C$ ，來自期望輸入大小為 $(N, C, +)$ 的輸入。
eps (float) – 為了數值穩定性新增到分母中的值。預設值：1e-5
momentum (Optional[float]) – 用於 running_mean 和 running_var 計算的值。可以設定為 None 以進行累積移動平均（即簡單平均）。預設值：0.1
affine (bool) – 一個布林值，如果設定為 True，則該模組具有可學習的仿射引數。預設值：True
track_running_stats (bool) – 一個布林值，當設定為 True 時，此模組會跟蹤執行均值和方差，當設定為 False 時，此模組不會跟蹤此類統計量，並將統計量緩衝區 running_mean 和 running_var 初始化為 None。當這些緩衝區為 None 時，此模組在訓練和評估模式下始終使用 batch 統計量。預設值：True
process_group (Optional[Any]) – 統計量的同步在每個程序組內部單獨進行。預設行為是在整個 world（所有程序）之間進行同步。

形狀

輸入： $(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)

classmethod convert_sync_batchnorm(module, process_group=None)[source]#

將模型中所有 BatchNorm*D 層轉換為 torch.nn.SyncBatchNorm 層。

引數

module (nn.Module) – 包含一個或多個 BatchNorm*D 層的模組。
process_group (optional) – 用於限定同步範圍的程序組，預設為整個 world。

返回

轉換後帶有 torch.nn.SyncBatchNorm 層的原始 module。如果原始 module 本身就是一個 BatchNorm*D 層，則會返回一個新的 torch.nn.SyncBatchNorm 層物件。

示例

>>> # Network with nn.BatchNorm layer
>>> module = torch.nn.Sequential(
>>>            torch.nn.Linear(20, 100),
>>>            torch.nn.BatchNorm1d(100),
>>>          ).cuda()
>>> # 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]
>>> sync_bn_module = torch.nn.SyncBatchNorm.convert_sync_batchnorm(module, process_group)

forward(input)[source]#

執行前向傳播。

返回型別: 張量

SyncBatchNorm#

文件

教程

資源