(self, input, weight, bias, running_mean, running_var)
| 403 | class SyncBatchNormFunction(flow.autograd.Function): |
| 404 | @staticmethod |
| 405 | def forward(self, input, weight, bias, running_mean, running_var): |
| 406 | assert input.is_local, "SyncBatchNorm does not support global tensor as input." |
| 407 | |
| 408 | if not input.is_contiguous(): |
| 409 | input = input.contiguous() |
| 410 | if weight is not None: |
| 411 | weight = weight.contiguous() |
| 412 | |
| 413 | size = int(input.numel() // input.size(1)) |
| 414 | if size == 1 and global_world_size < 2: |
| 415 | raise ValueError( |
| 416 | "Expected more than 1 value per channel when training, got input size {}".format( |
| 417 | size |
| 418 | ) |
| 419 | ) |
| 420 | |
| 421 | num_channels = input.shape[global_axis] |
| 422 | if input.numel() > 0: |
| 423 | # calculate mean/invstd for input. |
| 424 | mean, invstd = flow._C.batch_norm_stats(input, global_axis, global_eps) |
| 425 | |
| 426 | count = flow.full( |
| 427 | (1,), |
| 428 | input.numel() // input.size(global_axis), |
| 429 | dtype=mean.dtype, |
| 430 | device=mean.device, |
| 431 | ) |
| 432 | |
| 433 | # C, C, 1 -> (2C + 1) |
| 434 | combined = flow.cat([mean, invstd, count], dim=0) |
| 435 | else: |
| 436 | # for empty input, set stats and the count to zero. The stats with |
| 437 | # zero count will be filtered out later when computing global mean |
| 438 | # & invstd, but they still needs to participate the all_gather |
| 439 | # collective communication to unblock other peer processes. |
| 440 | combined = flow.zeros( |
| 441 | 2 * num_channels + 1, dtype=input.dtype, device=input.device |
| 442 | ) |
| 443 | |
| 444 | # Use allgather instead of allreduce because count could be different across |
| 445 | # ranks, simple all reduce op can not give correct results. |
| 446 | # batch_norm_gather_stats_with_counts calculates global mean & invstd based on |
| 447 | # all gathered mean, invstd and count. |
| 448 | # world_size * (2C + 1) |
| 449 | combined_size = combined.numel() |
| 450 | combined_flat = flow.empty( |
| 451 | global_world_size, |
| 452 | combined_size, |
| 453 | dtype=combined.dtype, |
| 454 | device=combined.device, |
| 455 | ) |
| 456 | flow.comm.all_gather_into_tensor(combined_flat, combined) |
| 457 | # world_size * (2C + 1) -> world_size * C, world_size * C, world_size * 1 |
| 458 | mean_all, invstd_all, count_all = flow.split(combined_flat, num_channels, dim=1) |
| 459 | |
| 460 | # remove stats from empty inputs |
| 461 | mask = count_all.squeeze(-1) >= 1 |
| 462 | count_all = count_all[mask] |
nothing calls this directly
no test coverage detected