语义分割—交叉熵损失函数

from dataloaders import utils
criterion = utils.cross_entropy2d
 
import torch.nn as nn
def cross_entropy2d(logit, target, ignore_index=255, weight=None, size_average=True, batch_average=True):
    n, c, h, w = logit.size()
    # logit = logit.permute(0, 2, 3, 1)
    target = target.squeeze(1)
    if weight is None:
        criterion = nn.CrossEntropyLoss(weight=weight, ignore_index=ignore_index, size_average=False)
    else:
        criterion = nn.CrossEntropyLoss(weight=torch.from_numpy(np.array(weight)).float().cuda(), ignore_index=ignore_index, size_average=False)
    loss = criterion(logit, target.long())
 
    if size_average:
        loss /= (h * w)
 
    if batch_average:
        loss /= n
 
    return loss

最近发现一个问题， 就是语义分割target这一步有的代码会将target.squeeze,  有的不需要呢   我测试了一下  如果标签单通道的灰度图的话  Image.open()读取是二维， 然后经过DataLoader，