python - Torchvision 的 CUDA 错误设备端断言错误

Question

我遇到了 Pytorch 和 CUDA 的问题，这似乎很常见，在使用 CUDA 构建神经网络时出现“RuntimeError：CUDA 错误：设备端断言已触发”，因为模型的输出与用于计算的标签不兼容失利。我认为这是因为模型的输出大小为 [n, 1] 而不是 [n] ，其中 n 是批量大小。

使用我自己的架构，我可以通过在前向传递函数中使用 .squeeze() 来解决此问题，以使输出大小为 [n]。现在我正在使用来自 torchvision 的预训练 ResNet 模型，但我无法实现同样的目标。我尝试为最后一层定义自己的模块，但没有奏效：

class final_layers(nn.Module):
    def __init__(self):
        super(final_layers, self).__init__()
        self.layers = nn.Sequential(
            nn.Linear(512, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.layers(x)
        return x.squeeze()

model = models.resnet18(pretrained=True)
model.fc = final_layers()
optimizer = Adam(model.parameters(), lr=0.0003)
criterion = BCELoss()
if torch.cuda.is_available():
    model = model.cuda()
    criterion = criterion.cuda()

该模型在带有 CPU 的 Colab 中运行良好，但在使用 GPU 时出现以下错误：

RuntimeError: CUDA error: device-side assert triggered

任何帮助将不胜感激 ：）

标签： pythonpytorchconv-neural-network