PyTorch BatchNorm2d报错：running_mean元素数不符的技术求助

错误根源 #

你犯了一个基础错误：BatchNorm2d的初始化参数必须是输入张量的通道数C，而不是高度H。

看你报错的代码：

import torch
        
n = 32 # N = Batch size
c = 1 # C = Channels
h = 64 # H = Height
w = 512 # W = Width
        
torch.nn.BatchNorm2d(h)(torch.rand(n,c,h,w))

这里你把高度h=64传给了BatchNorm2d，但输入张量的通道数是c=1，两者维度不匹配，直接触发错误。

为什么NHWC格式的代码能运行？ #

你提供的第二个代码能运行完全是巧合：

import torch

n = 32 # N = Batch size
c = 1 # C = Channels
h = 64 # H = Height
w = 512 # W = Width

x = torch.rand(n,h,w,c)

x = torch.nn.BatchNorm2d(h)(x)

你把h=64传给BatchNorm2d后，输入的NHWC张量中第二个维度是h=64，BatchNorm2d会错误地将这个维度当成通道数处理，刚好参数匹配。但这是完全的误用——它实际是对高度维度做归一化，根本不是你想要的通道维度归一化效果。

正确解决方法 #

方法1：按NCHW格式正确使用BatchNorm2d

初始化BatchNorm2d时传入通道数c，输入保持NCHW格式：

import torch

n = 32  # 批量大小
c = 1   # 通道数
h = 64  # 高度
w = 512 # 宽度

# 输入为NCHW格式，BatchNorm2d指定通道数c
x = torch.rand(n, c, h, w)
x = torch.nn.BatchNorm2d(c)(x)

方法2：如果输入是NHWC格式，先转成NCHW

如果你的数据是NHWC格式，先通过permute重排维度为NCHW，再使用BatchNorm2d：

import torch

n = 32  # 批量大小
c = 1   # 通道数
h = 64  # 高度
w = 512 # 宽度

# 输入为NHWC格式
x_nhwc = torch.rand(n, h, w, c)
# 转换为NCHW格式：N H W C → N C H W
x_nchw = x_nhwc.permute(0, 3, 1, 2)
# 传入通道数c初始化BatchNorm2d
x = torch.nn.BatchNorm2d(c)(x_nchw)

关键提醒 #

PyTorch的BatchNorm2d是针对通道维度做归一化，必须严格遵循NCHW输入格式，且初始化参数必须是通道数，这是官方明确规定的要求。

内容的提问来源于stack exchange，提问作者WhatAMesh