图像Captioning微调报错：running_mean应含1个元素而非256

我正在使用一个图像Captioning教程，针对自定义数据集进行模型微调。目前仅使用1张图像和1条标注训练，此前已解决将BatchNorm1d替换为InstanceNorm1d的问题，现在运行代码出现错误，相关代码如下：

编码器代码 #

class EncoderCNN(nn.Module):
    def __init__(self, embed_size):
        super(EncoderCNN, self).__init__()
        resnet = models.resnet152(pretrained=True)
        modules = list(resnet.children())[:-1]
        self.resnet = nn.Sequential(*modules)
        self.linear = nn.Linear(resnet.fc.in_features, embed_size)
        #self.bn = nn.BatchNorm1d(embed_size, momentum=0.01)
        self.bn = nn.InstanceNorm1d(256, track_running_stats=True)

    def forward(self, images):
        with torch.no_grad():
            features = self.resnet(images)
        features = features.reshape(features.size(0), -1)
        features = self.bn(self.linear(features))
        return features

数据加载器代码 #

class CustomDataset(Dataset):
    def __init__(self, image_folder, caption_folder, transform=None):
        self.image_folder = image_folder
        self.caption_folder = caption_folder
        self.transform = transform

    def __len__(self):
        return 256
        #return len(os.listdir(self.image_folder))
    
    def __getitem__(self, idx):
        #img_name = os.path.join(self.image_folder, str(idx) + '.jpg')
        img_name = os.path.join(self.image_folder, "0" + '.jpg')
        img = Image.open(img_name)
        

        transform = transforms.Compose([ 
            transforms.Resize((256,256)),
            transforms.RandomHorizontalFlip(), 
            transforms.ToTensor(), 
            transforms.Normalize((0.485, 0.456, 0.406), 
                                (0.229, 0.224, 0.225))])

        img = transform(img)

        def tokenize(caption, vocabulary):
            words = caption.split()
            tokens = [vocabulary(word) for word in words]
            return tokens

        tokenized_captions = tokenize(caption, vocab)

        print("Image shape:", img.shape)

        return img, torch.tensor(tokenized_captions)

dataset = CustomDataset(image_folder='images/',
                        caption_folder='captions/',
                        transform=transforms.ToTensor())
dataloader = DataLoader(dataset, batch_size=1, shuffle=True, drop_last=False)

模型加载代码 #

class FineTuneModel(nn.Module):
    def __init__(self, embed_size, hidden_size, vocab_size, num_layers):
        super(FineTuneModel, self).__init__()
        self.encoder = EncoderCNN(embed_size)
        self.decoder = DecoderRNN(embed_size,
                                  hidden_size,
                                  vocab_size,
                                  num_layers,
                                  max_seq_length=20)
    def forward(self, images, captions, lengths):
        features = self.encoder(images)
        outputs = self.decoder(features, captions, lengths)
        return outputs

1. InstanceNorm1d维度不匹配问题 #

EncoderCNN中，self.bn = nn.InstanceNorm1d(256, track_running_stats=True) 硬编码了特征维度为256，但self.linear的输出维度是传入的embed_size，若embed_size不等于256会直接导致维度不匹配错误。需将InstanceNorm1d的特征维度改为embed_size：

self.bn = nn.InstanceNorm1d(embed_size, track_running_stats=True)

2. caption变量未定义 #

在__getitem__方法中，调用tokenize(caption, vocab)时caption变量未被赋值。需从caption_folder中读取对应图像的标注文件，比如对应图像0.jpg的标注文件为0.txt，添加读取逻辑：

# 在tokenize函数定义前添加标注读取代码
caption_path = os.path.join(self.caption_folder, "0.txt")
with open(caption_path, 'r', encoding='utf-8') as f:
    caption = f.read().strip()

3. 数据集长度与实际数据不符 #

__len__返回256，但实际仅存在1张图像，会导致DataLoader重复读取同一张图，不符合数据集逻辑，应改为返回实际图像数量：

def __len__(self):
    return len(os.listdir(self.image_folder))

4. 变长文本序列处理问题 #

即使batch_size=1，后续调整batch_size大于1时，不同caption的长度不一致会触发默认collate_fn报错。需自定义collate_fn处理变长序列：

def collate_fn(data):
    # 按caption长度降序排序
    data.sort(key=lambda x: len(x[1]), reverse=True)
    images, captions = zip(*data)
    images = torch.stack(images, 0)
    
    # 生成长度列表和填充后的目标张量
    lengths = [len(cap) for cap in captions]
    targets = torch.zeros(len(captions), max(lengths)).long()
    for i, cap in enumerate(captions):
        end = lengths[i]
        targets[i, :end] = cap[:end]
    return images, targets, lengths

# 初始化DataLoader时传入自定义collate_fn
dataloader = DataLoader(dataset, batch_size=1, shuffle=True, drop_last=False, collate_fn=collate_fn)

5. 冗余的transform参数 #

Dataset初始化时传入了transform=transforms.ToTensor()，但__getitem__内又重新定义并应用transform，导致重复转换。可将transform统一移至__init__中：

# 修改Dataset的__init__和__getitem__
def __init__(self, image_folder, caption_folder, transform=None):
    self.image_folder = image_folder
    self.caption_folder = caption_folder
    # 若未传入transform则使用默认预处理逻辑
    if transform is None:
        self.transform = transforms.Compose([ 
            transforms.Resize((256,256)),
            transforms.RandomHorizontalFlip(), 
            transforms.ToTensor(), 
            transforms.Normalize((0.485, 0.456, 0.406), 
                                (0.229, 0.224, 0.225))])
    else:
        self.transform = transform

def __getitem__(self, idx):
    # ... 图像读取逻辑 ...
    img = self.transform(img)
    # ... 其余逻辑 ...

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