联邦学习vgg修改最后一个全连接层报错（默认是1000个输出，我的改为11就报错，改为200就没问题）

Question

联邦学习代码，我想实现vgg16训练自己的数据集，我的数据集28*28的灰度图，共11个类别，只要把
model.classifier[6] = nn.Linear(4096,1000,bias=True)
改为
model.classifier[6] = nn.Linear(4096,11,bias=True)
就会报错（resnet50 和alexnet就没这个问题）

Assertion `t >= 0 && t < n_classes` failed.
  torch.ones_like(out, memory_format=torch.preserve_format)
RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

而我如果改成：
model.classifier[6] = nn.Linear(4096,200,bias=True)就不报错
各位老哥帮忙解决一下
下面的是我有问题的vgg模型

VGG(
  (features): Sequential(
    (0): Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.8, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.8, inplace=False)
    (6): Linear(in_features=4096, out_features=11, bias=True)
  )
)

Answer 1

引自免费微信小程序：皆我百晓生

针对您遇到的问题，您在尝试将VGG16模型的最后一个全连接层（model.classifier[6]）的输出维度从1000修改为11时，遇到了特定的运行时错误。根据您提供的信息，错误提示包括一个设备端断言触发（CUDA error: device-side assert triggered）以及与类别的数量有关的断言失败（Assertion t >= 0 && t < n_classes failed.）。您提到ResNet50和AlexNet在同样的修改操作下没有问题，这表明问题可能与VGG16模型结构或您的训练设置有关。

以下是针对您描述的问题，可能的排查步骤及解决方案：

1. 检查数据标签： 确保您的数据集标签（ground truth labels）都是有效的整数索引，并且范围在0到10之间（对应于您数据集的11个类别）。如果存在标签超出此范围的情况，可能会触发上述错误。可以使用以下代码检查数据集中标签的有效性：

   import numpy as np
   
   # 假设您的标签存储在一个numpy数组labels中
   labels = ...  # 获取您的标签数据
   
   invalid_labels = np.where((labels < 0) | (labels >= 11))[0]
   if len(invalid_labels) > 0:
       print(f"发现{len(invalid_labels)}个无效标签：{invalid_labels}")
   else:
       print("所有标签有效")
   

2. 检查损失函数： 确认您使用的损失函数是否适用于多分类任务且与您的修改后的输出维度兼容。对于多分类问题，常用的损失函数如CrossEntropyLoss应能正确处理不同数量的类别。查看您的训练代码，确认损失函数的定义和使用：

   from torch.nn import CrossEntropyLoss
   
   criterion = CrossEntropyLoss()
   

   如果您使用了其他自定义损失函数或参数，请确保它们与新的输出维度相匹配。

3. 检查优化器状态： 如果您在修改模型结构前已经对模型进行过预训练或微调，那么优化器可能保存了与原模型结构相关的状态。在修改模型后，需要重新初始化优化器，以避免潜在的不匹配问题：

   optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=weight_decay)
   

4. 验证模型在CPU上的行为： 将模型和数据转移到CPU上进行训练，以排除可能的CUDA相关问题。如果在CPU上训练没有问题，那可能是CUDA环境或版本导致的异常。您可以尝试更新PyTorch版本或检查CUDA安装。

5. 逐步调整输出维度： 为了进一步定位问题，可以尝试逐步减小输出维度（如从1000到500、250等），观察在哪个点开始出现错误。这有助于确定是否存在特定输出维度值导致的问题。

6. 启用CUDA_LAUNCH_BLOCKING=1： 根据错误提示，可以尝试在运行训练脚本时设置环境变量CUDA_LAUNCH_BLOCKING=1，以便在发生CUDA错误时获取更准确的堆栈跟踪，有助于定位问题所在：

   export CUDA_LAUNCH_BLOCKING=1
   python your_training_script.py
   

如果以上步骤未能解决问题，请分享更多关于训练设置（如损失函数、优化器、数据加载等）和完整的错误堆栈信息，以便进行更深入的分析。