Upload 16 files

model/__init__.py

@@ -0,0 +1 @@
+#

model/resnet_backbone.py

@@ -0,0 +1,216 @@
+import math
+import torch.nn as nn
+
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding
+
+    Args:
+        in_planes (int): 输入通道数
+        out_planes (int): 输出通道数
+        stride (int): 卷积步幅，默认为 1
+        groups (int): 卷积分组数，默认为 1
+        dilation (int): 卷积扩张系数，默认为 1
+
+    Returns:
+        nn.Conv2d: 定义的 3x3 卷积层
+    """
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution
+    Args:
+        in_planes (int): 输入通道数
+        out_planes (int): 输出通道数
+        stride (int): 步幅，默认为 1
+
+    Returns:
+        nn.Conv2d: 定义的 1x1 卷积层
+    """
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        """
+                        初始化 Bottleneck 模块
+
+                        Args:
+                            inplanes (int): 输入通道数
+                            planes (int): 输出通道数（卷积层的通道数）
+                            stride (int): 卷积步幅，默认为 1
+                            downsample (nn.Module, optional): 用于调整输入尺寸的层（例如：当步幅不为 1 时，需要下采样）
+                            groups (int): 分组卷积的数量，默认为 1
+                            base_width (int): 基础宽度，影响每层的宽度
+                            dilation (int): 卷积扩张系数，默认为 1
+                            norm_layer (nn.Module, optional): 归一化层类型，默认为 `nn.BatchNorm2d`
+        """
+        super(Bottleneck, self).__init__()
+
+        # 如果没有提供 norm_layer，则使用 BatchNorm2d
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+
+        # 计算宽度，groups 会影响每个卷积层的通道数
+        width = int(planes * (base_width / 64.)) * groups
+
+        # 第一层 1x1 卷积，作用是降维
+        self.conv1 = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)  # 归一化层
+        # 第二层 3x3 卷积，作用是进行空间下采样
+        self.conv2 = conv3x3(width, width, stride, groups, dilation)
+        self.bn2 = norm_layer(width)  # 归一化层
+        # 第三层 1x1 卷积，作用是升维
+        self.conv3 = conv1x1(width, planes * self.expansion)
+        self.bn3 = norm_layer(planes * self.expansion)  # 归一化层
+
+        # 激活函数 ReLU
+        self.relu = nn.ReLU(inplace=True)
+
+        # 下采样层，默认为 None
+        self.downsample = downsample
+        # 存储步幅
+        self.stride = stride
+
+    def forward(self, x):
+        """
+                        正向传播
+
+                        Args:
+                            x (Tensor): 输入张量
+
+                        Returns:
+                            Tensor: 输出张量
+        """
+        identity = x  # 存储输入张量，用于跳跃连接（residual connection）
+
+        # 通过第一个 1x1 卷积层
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        # 通过第二个 3x3 卷积层
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        # 通过第三个 1x1 卷积层
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        # 跳跃连接：将输入（identity）与输出相加
+        out += identity
+
+        # 激活函数（ReLU）再次作用在输出上
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, layers, num_classes=1000):
+        super(ResNet, self).__init__()
+
+        # 设置初始通道数，供残差模块使用（后续根据 expansion 自动更新）
+        self.inplanes = 64
+
+        # 输入通道数为 3，输出为 64，7x7 大卷积核，步长为 2，padding 为 3，不使用 bias,为一个CBR模块
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+
+        # 最大池化层，3x3 核，步长为 2，无 padding，向上取整模式
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=0, ceil_mode=True)  # change
+
+        # 构建四个 stage 的残差网络，每个 stage 包含 layers[i] 个残差块
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+
+        # 平均池化，用于分类场景下整合空间信息（此实现中未启用）
+        self.avgpool = nn.AvgPool2d(7)
+        # 分类输出全连接层：输入通道为最后一层的输出通道（含 expansion），输出为类别数
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        # 模块初始化
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                # He 初始化：N(0, sqrt(2/n))，适合 ReLU 激活函数
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                # BN 权重初始化为 1，偏置初始化为 0
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        """
+                构建一个残差层（stage），由多个残差块 block 组成。
+                - block: 残差块类型（如 Bottleneck 或 BasicBlock）
+                - planes: 当前 stage 的基础输出通道数
+                - blocks: 当前 stage 包含的残差块数量
+                - stride: 第一个 block 的步长（用于空间下采样）
+        """
+        downsample = None
+        # 如果输入通道与输出通道不一致，或 stride ≠ 1，则需要下采样对齐
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+        layers = []
+        # 第一个 block：可能包含下采样和通道数调整
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        # 更新当前输入通道数为新的输出通道数
+        self.inplanes = planes * block.expansion
+
+        # 其余 block：不下采样，通道保持一致
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        # print("输入数据：", x.shape)
+        """
+                定义前向传播过程，输出多个阶段的特征图（适合用于下游任务如 FPN、检测、分割等）
+        """
+        x = self.conv1(x)
+        x = self.bn1(x)
+        feat1 = self.relu(x)  # 输出特征图
+        # print("feat1:", feat1.shape)
+
+        x = self.maxpool(feat1)
+        # print("池化层：", x.shape)
+
+        feat2 = self.layer1(x)  # stage 1 输出
+        feat3 = self.layer2(feat2)  # stage 2 输出
+        feat4 = self.layer3(feat3)   # stage 3 输出
+        feat5 = self.layer4(feat4)   # stage 4 输出
+
+        # 返回所有层的中间特征图（而非分类结果），常用于特征提取任务
+        return [feat1, feat2, feat3, feat4, feat5]
+
+
+def resnet50(**kwargs):
+    # 构建一个 ResNet-50 模型实例，使用 Bottleneck 结构
+    # 每个 stage 分别包含 3, 4, 6, 3 个残差块（符合 ResNet-50 配置）
+    # 额外参数通过 kwargs 传入 ResNet 构造函数（如 num_classes, input_shape 等）
+    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
+
+    # 删除全局平均池化层（avgpool）和全连接层（fc），以去除分类模块
+    # 模型将仅保留主干部分（stem + stage1~stage4），适合特征提取任务
+    del model.avgpool
+    del model.fc
+
+    # 返回去除分类头的 ResNet-50 模型，用于下游任务（如检测、分割）
+    return model

model/unet_resnet.py

@@ -0,0 +1,104 @@
+import torch
+import torch.nn as nn
+
+from model.resnet_backbone import resnet50
+
+# 定义一个 U-Net 解码模块（上采样模块）
+class unetUp(nn.Module):
+    def __init__(self, in_size, out_size):
+        """
+               构造函数
+               参数：
+               - in_size: 输入通道数（等于上采样特征图通道数 + 跳跃连接特征图通道数）
+               - out_size: 输出通道数（经过卷积后的输出通道数）
+        """
+        super(unetUp, self).__init__()
+        # 第一层卷积：输入通道为 in_size，输出通道为 out_size，卷积核大小 3x3，padding=1 保持尺寸不变
+        self.conv1 = nn.Conv2d(in_size, out_size, kernel_size=3, padding=1)
+        # 第二层卷积：输入通道和输出通道均为 out_size，进一步特征提取
+        self.conv2 = nn.Conv2d(out_size, out_size, kernel_size=3, padding=1)
+        # 上采样层：使用双线性插值将特征图放大两倍（scale_factor=2）
+        self.up = nn.UpsamplingBilinear2d(scale_factor=2)
+        # ReLU 激活函数，inplace=True 表示原地操作，节省内存
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, inputs1, inputs2):
+        """
+               前向传播
+               参数：
+               - inputs1: 编码器部分通过跳跃连接传来的特征图（高分辨率）
+               - inputs2: 来自上一级解码器的输出（低分辨率，需要上采样）
+        """
+        # 先对 inputs2 进行上采样，然后与 inputs1 在通道维度上进行拼接
+        # 拼接后的通道数为 in_size
+        outputs = torch.cat([inputs1, self.up(inputs2)], 1)
+        # 第一次卷积 + ReLU
+        outputs = self.conv1(outputs)
+        outputs = self.relu(outputs)
+        # 第二次卷积 + ReLU
+        outputs = self.conv2(outputs)
+        outputs = self.relu(outputs)
+        # 返回该解码模块的输出特征图
+        return outputs
+
+
+# 定义 U-Net 主体结构
+class Unet(nn.Module):
+    def __init__(self, num_classes=21):
+        """
+                构造函数
+                参数：
+                - num_classes: 最终分类的类别数（用于语义分割任务）
+        """
+        super(Unet, self).__init__()
+
+        # 使用 ResNet50 作为编码器，提取多尺度特征
+        self.resnet = resnet50()  # 假设 resnet50() 返回 5 层特征：feat1 ~ feat5
+        # 编码器输出的特征通道数（每层输出的特征图通道数）
+        in_filters = [192, 512, 1024, 3072]   # 通常是低到高的顺序（层级）
+        # 解码器每一层希望恢复到的输出通道数
+        out_filters = [64, 128, 256, 512]   # 解码层输出通道数逐步降低
+
+        # 定义 4 层上采样模块（从深到浅）
+        # 每层通过双线性插值上采样 + 拼接 + 两次卷积 + ReLU
+        self.up_concat4 = unetUp(in_filters[3], out_filters[3])  # 最深层输出
+        self.up_concat3 = unetUp(in_filters[2], out_filters[2])  # 次深层
+        self.up_concat2 = unetUp(in_filters[1], out_filters[1])  # 中间层
+        self.up_concat1 = unetUp(in_filters[0], out_filters[0])  # 最浅层
+
+        # 对最后一层解码器的输出再上采样一倍并做两次卷积处理，进一步提升分辨率
+        self.up_conv = nn.Sequential(
+                nn.UpsamplingBilinear2d(scale_factor=2),   # 上采样
+                nn.Conv2d(out_filters[0], out_filters[0], kernel_size=3, padding=1),  # 卷积
+                nn.ReLU(),
+                nn.Conv2d(out_filters[0], out_filters[0], kernel_size=3, padding=1),  # 卷积
+                nn.ReLU(),
+            )
+        # 最后一层卷积，用于生成最终每类像素的得分图（通道数=num_classes）
+        self.final = nn.Conv2d(out_filters[0], num_classes, 1)
+
+    def forward(self, inputs):
+        """
+                前向传播
+                参数：
+                - inputs: 输入图像（通常为 RGB 图像，形状为 B×3×H×W）
+
+                返回：
+                - final: 每个像素点在 num_classes 个类别上的预测（未Softmax）
+        """
+        # 编码器提取五层特征图（假设顺序为由浅到深）
+        [feat1, feat2, feat3, feat4, feat5] = self.resnet.forward(inputs)
+
+        # 解码过程：由深到浅依次上采样 + 融合跳跃连接特征
+        up4 = self.up_concat4(feat4, feat5)  # 使用 feat5 上采样后与 feat4 融合
+        up3 = self.up_concat3(feat3, up4)   # 再将上一层上采样后与 feat3 融合
+        up2 = self.up_concat2(feat2, up3)   # 同理
+        up1 = self.up_concat1(feat1, up2)   # 最浅层融合
+
+        # 可选的最终进一步上采样（恢复到输入图尺寸）
+        if self.up_conv != None:
+            up1 = self.up_conv(up1)
+
+        # 最终生成每类的预测图
+        final = self.final(up1)
+        return final

model/unet_training.py

@@ -0,0 +1,199 @@
+import math
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def CE_Loss(inputs, target, cls_weights, num_classes=21):
+    n, c, h, w = inputs.size()  # 输入大小，n=batch_size，c=类别数，h=高度，w=宽度
+    nt, ht, wt = target.size()  # 目标大小，nt=batch_size，ht=目标图像的高度，wt=目标图像的宽度
+
+    # 如果输入和目标的大小不一致，则进行插值操作
+    if h != ht and w != wt:
+        inputs = F.interpolate(inputs, size=(ht, wt), mode="bilinear", align_corners=True)
+
+    # 将输入和目标调整为合适的形状
+    temp_inputs = inputs.transpose(1, 2).transpose(2, 3).contiguous().view(-1, c)  # 调整为 [n*h*w, c]
+    # temp_inputs = inputs.view(-1, c)  # 变为 [n*h*w, c]
+    temp_target = target.view(-1)  # 展平目标标签为 [n*h*w]
+
+    # 计算交叉熵损失
+    CE_loss = nn.CrossEntropyLoss(weight=cls_weights, ignore_index=num_classes)(temp_inputs, temp_target)
+    return CE_loss
+
+
+"""
+这段代码实现了 Focal Loss，这是一个改进版的交叉熵损失，
+专门用于处理类别不平衡问题。它通过对 困难样本（即模型不确定的样本）
+施加更高的权重来增强模型对这些样本的关注，尤其是在训练过程中遇到类不平衡时。
+"""
+def Focal_Loss(inputs, target, cls_weights, num_classes=21, alpha=0.5, gamma=2):
+    n, c, h, w = inputs.size()  # 输入的尺寸，n=batch_size，c=类别数，h=高度，w=宽度
+    nt, ht, wt = target.size()  # 目标标签的尺寸，nt=batch_size，ht=目标高度，wt=目标宽度
+
+    # 如果输入和目标的尺寸不一致，则进行插值调整
+    if h != ht and w != wt:
+        inputs = F.interpolate(inputs, size=(ht, wt), mode="bilinear", align_corners=True)
+
+    # 展平输入和目标，使其适应交叉熵损失函数的要求
+    temp_inputs = inputs.transpose(1, 2).transpose(2, 3).contiguous().view(-1, c)
+    # temp_inputs = inputs.view(-1, c)  # 输入变为 [n*h*w, c]
+    temp_target = target.view(-1)  # 目标变为 [n*h*w]
+
+    # 计算交叉熵损失，使用 `reduction='none'` 返回每个像素的损失
+    logpt = -nn.CrossEntropyLoss(weight=cls_weights, ignore_index=num_classes, reduction='none')(temp_inputs,
+                                                                                                 temp_target)
+    pt = torch.exp(logpt)  # 计算每个像素属于真实类别的概率
+
+    # 如果 alpha 存在，则乘上 alpha 来调节正负样本的权重
+    if alpha is not None:
+        logpt *= alpha
+
+    # 计算 Focal Loss
+    loss = -((1 - pt) ** gamma) * logpt
+
+    # 对所有像素的损失求平均
+    loss = loss.mean()
+    return loss
+
+
+"""
+这段代码实现了 Dice Loss，它是用于计算图像分割任务中的性能度量指标，
+尤其适用于数据类别不平衡的场景。Dice 系数衡量的是预测结果和真实标签之间的重叠程度，
+通常用于语义分割任务中，特别是在医学图像分割等领域。
+"""
+def Dice_loss(inputs, target, beta=1, smooth=1e-5):
+    n, c, h, w = inputs.size()  # 输入的尺寸：n=batch_size, c=类别数, h=高度, w=宽度
+    nt, ht, wt, ct = target.size()  # 目标标签的尺寸：nt=batch_size, ht=高度, wt=宽度, ct=类别数
+
+    # 如果输入和目标的尺寸不一致，进行插值
+    if h != ht and w != wt:
+        inputs = F.interpolate(inputs, size=(ht, wt), mode="bilinear", align_corners=True)
+
+    # 对输入和目标进行转换，保证计算时的一致性
+    temp_inputs = torch.softmax(inputs.transpose(1, 2).transpose(2, 3).contiguous().view(n, -1, c), -1)
+    # temp_inputs = torch.softmax(inputs.permute(0, 2, 3, 1).contiguous().view(n, -1, c), -1)
+    temp_target = target.view(n, -1, ct)
+
+    # 计算真正例 (tp)，即真实类别和预测类别的交集
+    tp = torch.sum(temp_target[..., :-1] * temp_inputs, axis=[0, 1])
+    # 计算假正例 (fp)，即预测为类别而真实标签为背景
+    fp = torch.sum(temp_inputs, axis=[0, 1]) - tp
+    # 计算假负例 (fn)，即真实标签为类别而预测为背景
+    fn = torch.sum(temp_target[..., :-1], axis=[0, 1]) - tp
+
+    # 计算 Dice 系数，beta 控制对假负样本的惩罚程度
+    score = ((1 + beta ** 2) * tp + smooth) / ((1 + beta ** 2) * tp + beta ** 2 * fn + fp + smooth)
+    # 计算 Dice Loss
+    dice_loss = 1 - torch.mean(score)
+    return dice_loss
+
+
+def weights_init(net, init_type='normal', init_gain=0.02):
+    def init_func(m):
+        classname = m.__class__.__name__  # 获取模块的类名
+        if hasattr(m, 'weight') and classname.find('Conv') != -1:  # 如果是卷积层
+            if init_type == 'normal':
+                torch.nn.init.normal_(m.weight.data, 0.0, init_gain)  # 正态分布初始化
+            elif init_type == 'xavier':
+                torch.nn.init.xavier_normal_(m.weight.data, gain=init_gain)  # Xavier初始化
+            elif init_type == 'kaiming':
+                torch.nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')  # Kaiming初始化
+            elif init_type == 'orthogonal':
+                torch.nn.init.orthogonal_(m.weight.data, gain=init_gain)  # 正交初始化
+            else:
+                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+        elif classname.find('BatchNorm2d') != -1:  # 如果是批量归一化层
+            torch.nn.init.normal_(m.weight.data, 1.0, 0.02)  # 正态分布初始化，均值1，标准差0.02
+            torch.nn.init.constant_(m.bias.data, 0.0)  # 偏置初始化为0
+
+    print('initialize network with %s type' % init_type)  # 打印初始化方法
+    net.apply(init_func)  # 将初始化函数应用到网络的每一层
+
+
+def get_lr_scheduler(lr_decay_type, lr, min_lr, total_iters, warmup_iters_ratio=0.05, warmup_lr_ratio=0.1,
+                     no_aug_iter_ratio=0.05, step_num=10):
+    # 根据学习率衰减类型（lr_decay_type）选择不同的学习率调度方式
+    def yolox_warm_cos_lr(lr, min_lr, total_iters, warmup_total_iters, warmup_lr_start, no_aug_iter, iters):
+        """
+                YOLOX特定的余弦衰减带热身策略的学习率调度函数。
+                根据当前的训练迭代次数 (iters) 来动态调整学习率。
+
+                参数：
+                lr: 当前的学习率
+                min_lr: 最小学习率
+                total_iters: 总的训练迭代次数
+                warmup_total_iters: 热身阶段的迭代次数
+                warmup_lr_start: 热身阶段开始时的学习率
+                no_aug_iter: 不进行数据增强的迭代次数
+                iters: 当前的训练迭代次数
+
+                返回：
+                调整后的学习率
+        """
+        if iters <= warmup_total_iters:
+            # 在热身阶段，使用一个二次函数来逐步增加学习率
+            lr = (lr - warmup_lr_start) * pow(iters / float(warmup_total_iters), 2) + warmup_lr_start
+        elif iters >= total_iters - no_aug_iter:
+            # 如果迭代数接近总迭代数，且接近不使用数据增强的迭代次数，则将学习率设置为最小学习率
+            lr = min_lr
+        else:
+            # 在其他阶段，使用余弦衰减函数来调整学习率
+            lr = min_lr + 0.5 * (lr - min_lr) * (
+                    1.0 + math.cos(
+                math.pi * (iters - warmup_total_iters) / (total_iters - warmup_total_iters - no_aug_iter))
+            )
+        return lr
+
+    def step_lr(lr, decay_rate, step_size, iters):
+        """
+                逐步衰减的学习率调度函数。
+
+                参数：
+                lr: 当前的学习率
+                decay_rate: 衰减率，用来控制每个step后学习率的减少比例
+                step_size: 每次衰减发生的步数
+                iters: 当前的训练迭代次数
+
+                返回：
+                调整后的学习率
+        """
+        if step_size < 1:
+            raise ValueError("step_size must above 1.")
+        # 计算经过了多少个衰减步
+        n = iters // step_size
+        # 按照衰减率调整学习率
+        out_lr = lr * decay_rate ** n
+        return out_lr
+
+    # 根据学习率衰减类型选择合适的调度函数
+    if lr_decay_type == "cos":
+        # 如果选择的是余弦衰减（cos），计算热身阶段的迭代次数和学习率
+        warmup_total_iters = min(max(warmup_iters_ratio * total_iters, 1), 3)
+        warmup_lr_start = max(warmup_lr_ratio * lr, 1e-6)
+        no_aug_iter = min(max(no_aug_iter_ratio * total_iters, 1), 15)
+        # 使用partial函数将固定参数传给yolox_warm_cos_lr
+        func = partial(yolox_warm_cos_lr, lr, min_lr, total_iters, warmup_total_iters, warmup_lr_start, no_aug_iter)
+    else:
+        # 否则，使用逐步衰减的策略
+        # 计算每次衰减的速率
+        decay_rate = (min_lr / lr) ** (1 / (step_num - 1))
+        # 计算每步的衰减步长
+        step_size = total_iters / step_num
+        # 使用partial函数将固定参数传给step_lr
+        func = partial(step_lr, lr, decay_rate, step_size)
+
+    # 返回选择好的学习率调度函数
+    return func
+
+
+def set_optimizer_lr(optimizer, lr_scheduler_func, epoch):
+    # 使用学习率调度函数 lr_scheduler_func，根据当前的 epoch 获取对应的学习率
+    lr = lr_scheduler_func(epoch)
+
+    # 遍历优化器的所有参数组（optimizer.param_groups 是一个包含所有参数组的列表）
+    for param_group in optimizer.param_groups:
+        # 将当前参数组的学习率（lr）更新为调度函数计算得到的学习率
+        param_group['lr'] = lr