dimension out of range [-1, 0], but got 1：UNet 的 CrossEntropyLoss 不要把类别维压平

原帖自定义损失先 sigmoid，再把输出和掩码都 view(-1)，然后传给交叉熵：

RuntimeError: dimension out of range (expected to be in range of [-1, 0], but got 1)

类别轴丢了，改 dim 数字不能补回来

原题已经在损失中把 logits 与掩码展平成向量。此时再对 dim=1 操作，张量只剩维度 0，自然超出范围。关键是保持“每个像素有两个类别分数”的布局，而不是把所有 dim 都改成 0。

这会把 [N,C,H,W] 压成一维，交叉熵所需的类别维随之丢失。原题网络末层有两个通道，修复重点是损失入口的布局和输出含义。

在自定义损失的 forward(logits, targets) 开头检查：

print("josh050608")
print("logits:", logits.shape, logits.dtype)
print("targets:", targets.shape, targets.dtype)
print("target values:", targets.unique()[:20])

若是两个互斥类别的逐像素分割，让网络最后返回 self.last(up4) 的原始 logits，删去输出处的 log_softmax，也删去损失中的 sigmoid 和两个 view(-1)。损失直接写：

import torch.nn.functional as F
# targets 是已经映射成 0/1 类别编号的掩码
loss = F.cross_entropy(logits, targets.long())

此时 logits 应为 [N,2,H,W]，targets 为 [N,H,W]。如果掩码是 [N,1,H,W]，确认这一维只是单通道后用 targets.squeeze(1)；如果是 RGB 颜色掩码，先按颜色到类别的映射转换，不能直接 .long() 当作类别。

若实际任务采用单输出通道二分类，应改为 [N,1,H,W] logits 和同形状浮点 0/1 掩码，使用 BCEWithLogitsLoss。这是另一个明确方案，不要混用两个约定。

先验证一个 batch：空间尺寸一致、标签范围合法、loss 为标量，并执行反向传播。再检查 logits.argmax(dim=1) 是否得到与掩码同尺寸的类别图。

两通道分割的完整最小验证

下面用一个 1×1 卷积模拟 UNet 最后输出：两张 8×8 图片，每个像素只有 0/1 一个类别编号。可在 CPU 直接运行，不依赖原题数据：

import torch
from torch import nn
torch.manual_seed(12)
head = nn.Conv2d(3, 2, kernel_size=1)
images = torch.randn(2, 3, 8, 8)
mask = torch.randint(0, 2, (2, 8, 8), dtype=torch.long)
logits = head(images)
assert logits.shape == (2, 2, 8, 8)
assert mask.shape == (2, 8, 8)
loss = nn.CrossEntropyLoss()(logits, mask)
assert loss.ndim == 0 and torch.isfinite(loss)
loss.backward()
assert head.weight.grad is not None
prediction = logits.argmax(dim=1)
assert prediction.shape == mask.shape
print("loss/prediction:", float(loss.detach()), prediction.shape)

把例子迁回原 UNet 时，logits 应直接来自最后的卷积，不要提前 sigmoid、softmax 或 flatten。掩码的 unique() 如果是 0 和 255，需要按标注定义把前景映射为 1；若 255 是忽略像素，则明确设置 ignore_index，不能把它也当作前景。转换为 long 只改存储类型，不改变标签含义。

选定一种设计后，连同预测阈值或 argmax 轴一起保持一致。验收时抽一张掩码可视化，确认像素编号与前景含义，随后检查完整 batch 和 N=1 的空间尺寸；loss 能计算只是第一步，类别轴正确才是这次修复的目标。

参考：原始求助帖；技术依据：CrossEntropyLoss（2.5）、BCEWithLogitsLoss（2.5）。