在训练模式下，DeepLabV3的输入应该是什么

Question

我试图在自定义数据集上训练deeplabv3_resnet50模型，但在尝试进行前向传递时得到错误ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 256, 1, 1])。以下最小示例会产生此错误：

import torch
import torchvision
model = torchvision.models.segmentation.deeplabv3_resnet50(weights="DEFAULT")
model.train()
batch_size = 1
nbr_of_channels = 3
img_height, img_width = (500, 500)
input = torch.rand((batch_size, nbr_of_channels, img_height, img_width))
model(input)

我一点也不明白。got input size torch.Size([1, 256, 1, 1])是什么意思，我应该做些什么不同的事情？

Answer 1

您得到的错误来自深层BatchNorm层：在骨干网的深层，特征图大小减少到1x1像素。因此，当批次大小只有一个时，BatchNorm无法计算特征图的std。

对于任何batch_size > 1，这将起作用：

batch_size = 2  # Need bigger batches for training
nbr_of_channels = 3
img_height, img_width = (500, 500)
input = torch.rand((batch_size, nbr_of_channels, img_height, img_width))
model(input)  # working with batch_size > 1