|
|
发表于 2025-4-24 16:04:32
|
显示全部楼层
# ====================== 3. 简化版YOLOv2模型 ======================
class SimpleYOLO(nn.Module):
def __init__(self, num_classes=10, anchors = torch.tensor([[1.3221, 1.7314]])):
#锚框的宽高比例（宽度是高度的1.3221倍）
#锚框:模型可以快速适应不同大小的目标，而不需要从零开始预测边界框的绝对尺寸
super().__init__()## 必须调用父类初始化
self.num_classes = num_classes# 类别数量（MNIST有10个数字）
self.anchors = torch.tensor(anchors)# 预定义的锚框尺寸（这里只有1种）
# 简化网络结构
self.features = nn.Sequential(
nn.Conv2d(3, 16, 3, padding=1),#输入3通道，输出16通道，3x3卷积，边缘填充1
nn.LeakyReLU(0.1),# 激活函数
nn.MaxPool2d(2, 2),#最大池化
nn.Conv2d(16, 32, 3, padding=1),#输入16通道，输出32通道，3x3卷积，边缘填充1
nn.LeakyReLU(0.1),# 激活函数
nn.MaxPool2d(2, 2),#最大池化
nn.Conv2d(32, 64, 3, padding=1),#输入32通道，输出64通道，3x3卷积，边缘填充1
nn.LeakyReLU(0.1),# 激活函数
nn.MaxPool2d(2, 2),#最大池化
nn.Conv2d(64, 128, 3, padding=1),#输入64通道，输出128通道，3x3卷积，边缘填充1
nn.LeakyReLU(0.1),# 激活函数
nn.MaxPool2d(2, 2)#最大池化
)
# 检测层 用于预测结果
self.detection = nn.Conv2d(128, len(anchors)*(5+num_classes), 1)
## 输出通道数 = 锚框数量 × (5个边界框参数 + 类别数量)
## 5个边界框参数：x_center, y_center, width, height, confidence
def forward(self, x):#向前传播的过程
x = self.features(x)## 通过特征提取层（输出形状：[batch, 128, 26, 26]）#经过四次池化，每次都减半，416变为26
output = self.detection(x) ##通过检测层（输出形状：[batch, 1*(5+10), 26, 26]
#调整输出形状
batch_size,_, h, w = output.shape #获取高度，批次大小，宽度
output = output.view(batch_size, len(self.anchors), 5+self.num_classes, h, w) #调整为：[batch, 锚框数量, (5参数+10类别), 高度, 宽度]
output = output.permute(0, 1, 3, 4, 2) ## 调整维度顺序，变为：[batch, 锚框数量, 高度, 宽度, (5参数+10类别)]
return output
# ====================== 4. 简化版损失函数 ======================
def calculate_iou(pred_boxes, target_boxes):
"""计算IoU"""
# 将坐标中心坐标 → 左上角+右下角 #预测框
pred_boxes = pred_boxes.clone() #克隆张量
pred_boxes[:, 0] = pred_boxes[:, 0] - pred_boxes[:, 2]/2 #左上角x坐标
pred_boxes[:, 1] = pred_boxes[:, 1] - pred_boxes[:, 3]/2 #左上角y坐标
pred_boxes[:, 2] = pred_boxes[:, 0] + pred_boxes[:, 2] #右上角x坐标
pred_boxes[:, 3] = pred_boxes[:, 1] + pred_boxes[:, 3] #右上角y坐标
#目标框
target_boxes = target_boxes.clone() #同上
target_boxes[:, 0] = target_boxes[:, 0] - target_boxes[:, 2]/2
target_boxes[:, 1] = target_boxes[:, 1] - target_boxes[:, 3]/2
target_boxes[:, 2] = target_boxes[:, 0] + target_boxes[:, 2]
target_boxes[:, 3] = target_boxes[:, 1] + target_boxes[:, 3]
# 计算交集区域 计算预测框和目标框的交集区域面积
inter_x1 = torch.max(pred_boxes[:, 0], target_boxes[:, 0])#torch.max()：逐元素比较，取较大值
inter_y1 = torch.max(pred_boxes[:, 1], target_boxes[:, 1])
inter_x2 = torch.min(pred_boxes[:, 2], target_boxes[:, 2])#torch.min()：逐元素比较，取较小值
inter_y2 = torch.min(pred_boxes[:, 3], target_boxes[:, 3])
inter_area = torch.clamp(inter_x2 - inter_x1, min=0) * torch.clamp(inter_y2 - inter_y1, min=0)
#交集区域的宽度 = inter_x2 - inter_x1，高度 = inter_y2 - inter_y1，面积 = 宽度 × 高度
# 计算并集面积
pred_area = (pred_boxes[:, 2] - pred_boxes[:, 0]) * (pred_boxes[:, 3] - pred_boxes[:, 1])
target_area = (target_boxes[:, 2] - target_boxes[:, 0]) * (target_boxes[:, 3] - target_boxes[:, 1])
union_area = pred_area + target_area - inter_area
return inter_area / union_area
|
|
发表于 2025-3-14 20:20:21
楼主