File size: 7,569 Bytes
c2b1b26 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 | """
改进的损失函数 - 更好地处理类别不平衡
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
class ImprovedFocalLoss(nn.Module):
"""
改进的Focal Loss，支持类别权重
"""
def __init__(self, alpha=None, gamma=2, ignore_index=255, reduction='mean', class_weights=None):
"""
Args:
alpha: 类别权重列表，如果为None则使用平衡权重
gamma: 聚焦参数
ignore_index: 忽略的索引
reduction: 降维方式
class_weights: 类别权重tensor，shape为(num_classes,)
"""
super().__init__()
if alpha is None:
alpha = [1.0, 1.0] # 默认平衡权重
if isinstance(alpha, (list, tuple)):
alpha = torch.tensor(alpha, dtype=torch.float32)
self.alpha = alpha
self.gamma = gamma
self.ignore_index = ignore_index
self.reduction = reduction
self.class_weights = class_weights
def forward(self, inputs, targets):
"""
Args:
inputs: 预测值 (N, C, H, W)
targets: 目标值 (N, H, W)
"""
# 处理多输出格式
if isinstance(inputs, dict):
inputs = inputs['out']
# 忽略指定索引
if self.ignore_index is not None:
mask = targets != self.ignore_index
targets = targets[mask]
inputs = inputs.permute(0, 2, 3, 1)[mask]
else:
inputs = inputs.permute(0, 2, 3, 1).contiguous().view(-1, inputs.size(1))
targets = targets.view(-1)
# 计算交叉熵
ce_loss = F.cross_entropy(inputs, targets, reduction='none', weight=self.class_weights)
pt = torch.exp(-ce_loss)
# 应用alpha权重
if self.alpha is not None:
if self.alpha.device != targets.device:
self.alpha = self.alpha.to(targets.device)
alpha_t = self.alpha[targets]
focal_loss = alpha_t * (1 - pt) ** self.gamma * ce_loss
else:
focal_loss = (1 - pt) ** self.gamma * ce_loss
if self.reduction == 'mean':
return focal_loss.mean()
elif self.reduction == 'sum':
return focal_loss.sum()
else:
return focal_loss
class WeightedDiceLoss(nn.Module):
"""
加权Dice Loss，给不同类别不同的权重
"""
def __init__(self, num_classes=2, smooth=1e-6, class_weights=None):
"""
Args:
num_classes: 类别数
smooth: 平滑系数
class_weights: 类别权重tensor，shape为(num_classes,)
"""
super().__init__()
self.num_classes = num_classes
self.smooth = smooth
self.class_weights = class_weights
def forward(self, inputs, targets):
"""
Args:
inputs: 预测值 (N, C, H, W)
targets: 目标值 (N, H, W)
"""
# 处理多输出格式
if isinstance(inputs, dict):
inputs = inputs['out']
# 将预测转换为概率
inputs = torch.softmax(inputs, dim=1)
# 创建one-hot编码
targets_one_hot = F.one_hot(targets, num_classes=self.num_classes)
targets_one_hot = targets_one_hot.permute(0, 3, 1, 2).float()
# 计算每个类别的dice系数
intersection = (inputs * targets_one_hot).sum(dim=(2, 3)) # (N, C)
union = inputs.sum(dim=(2, 3)) + targets_one_hot.sum(dim=(2, 3)) # (N, C)
dice_score = (2. * intersection + self.smooth) / (union + self.smooth) # (N, C)
# 计算每个类别的dice loss
dice_loss_per_class = 1 - dice_score # (N, C)
# 应用类别权重
if self.class_weights is not None:
if self.class_weights.device != dice_loss_per_class.device:
self.class_weights = self.class_weights.to(dice_loss_per_class.device)
# 对每个类别应用权重
weighted_dice_loss = dice_loss_per_class * self.class_weights.unsqueeze(0)
dice_loss = weighted_dice_loss.mean()
else:
dice_loss = dice_loss_per_class.mean()
return dice_loss
class ImprovedSeaweedSegmentationLoss(nn.Module):
"""
改进的浒苔分割损失函数
更好地处理类别不平衡，特别是背景像素远多于浒苔像素的情况
"""
def __init__(self, num_classes=2, focal_alpha=None, focal_gamma=2,
dice_weight=0.5, focal_weight=1.0,
background_weight=1.0, foreground_weight=2.0):
"""
Args:
num_classes: 类别数
focal_alpha: Focal Loss的alpha参数（类别权重）
focal_gamma: Focal Loss的gamma参数
dice_weight: Dice Loss的权重
focal_weight: Focal Loss的权重
background_weight: 背景类别的权重（类别0）
foreground_weight: 前景类别（浒苔）的权重（类别1）
"""
super().__init__()
# 创建类别权重
class_weights = torch.tensor([background_weight, foreground_weight], dtype=torch.float32)
# 创建Focal Loss的alpha权重
if focal_alpha is None:
# 默认给前景更高的权重
focal_alpha = [background_weight, foreground_weight]
self.focal_loss = ImprovedFocalLoss(
alpha=focal_alpha,
gamma=focal_gamma,
class_weights=class_weights
)
self.dice_loss = WeightedDiceLoss(
num_classes=num_classes,
class_weights=class_weights
)
self.dice_weight = dice_weight
self.focal_weight = focal_weight
self.num_classes = num_classes
def forward(self, inputs, targets):
"""计算组合损失"""
focal_loss = self.focal_loss(inputs, targets)
dice_loss = self.dice_loss(inputs, targets)
total_loss = self.focal_weight * focal_loss + self.dice_weight * dice_loss
return {
'total_loss': total_loss,
'focal_loss': focal_loss,
'dice_loss': dice_loss
}
def calculate_class_weights_from_dataset(dataset, num_classes=2):
"""
从数据集中计算类别权重
用于平衡类别不平衡问题
"""
print("正在计算类别权重...")
total_pixels = 0
class_counts = torch.zeros(num_classes)
for i in range(len(dataset)):
sample = dataset[i]
mask = sample['mask']
# 统计每个类别的像素数
for c in range(num_classes):
class_counts[c] += (mask == c).sum().item()
total_pixels += mask.numel()
# 计算权重：总像素数 / (类别数 * 该类别的像素数)
class_weights = total_pixels / (num_classes * class_counts + 1e-6)
# 归一化权重
class_weights = class_weights / class_weights.sum() * num_classes
print(f"类别像素统计: {class_counts.tolist()}")
print(f"类别权重: {class_weights.tolist()}")
return class_weights
|