from typing import Any
|
from src.utils.register import (DataCreator, ModuleCreator, OptimizerCreator, MetricorCreator, SchedulerCreator)
|
import torch
|
|
|
"""
|
- ExampleModel
|
模型的forward方法需包含训练和测试的前向传播逻辑
|
通过创建模型时传入mode参数来区分训练和评估模式，结合`self.training`属性来判断当前模式
|
|
-----forward-----
|
** 输入：
|
只能有一个输入参数inputs
|
inputs中需要包含用于前向传播的图片和对应的真值标签
|
推荐格式为：
|
{
|
"batch_inputs": tensor, # 输入数据，未标准化(即0~255)
|
"ground_truth": {
|
"image_id": int,
|
"image_state": str, # ["已标注", "未标注"]
|
"type": str, # ["train", "val", "test"]
|
"image_path": str,
|
"pLabel": list[int] 图像级标签
|
"gtBoxes": {
|
"boxes": list[list], # 框左上角和右下角点坐标xyxy
|
"labels": list[int] # 框对应的标签
|
}
|
"gtSegments": {
|
"segments": list[list[list]], # 分割掩码外轮廓点xy
|
"labels": list[int] #分割标签
|
}
|
}
|
}
|
|
** 输出：
|
通过is_training属性和创建模型时传入的mode类型判断当前模式
|
在训练模式下，直接返回tensor损失，或者返回一个字典 {"losses": tensor}
|
在评估模式下，返回一个字典，包含预测结果和真值标签，推荐格式为
|
{
|
"pred_output": Any, # 预测结果
|
"ground_truth": Any # 真值标签，可以直接将上文的ground_truth传出
|
}
|
"""
|
|
class ExampleDataset(torch.utils.data.Dataset):
|
"""模拟数据集"""
|
def __init__(self, data_path):
|
self.size = 10
|
self.data_path = data_path
|
|
def __len__(self):
|
return self.size
|
|
def __getitem__(self, idx):
|
# 模拟数据和标签
|
x = torch.randn(3, 64, 64)
|
gt = {
|
"imageId": idx,
|
"imagePath": f"{self.data_path}/image_{idx}.jpg",
|
"iLevelLabel": [1, 2], # 假设有3个类别(good, defect1, defect2)
|
"gtBoxes": {
|
"boxes": [[0.1, 0.1, 0.5, 0.5],], # 一个框
|
"labels": [1] # 框对应的标签
|
},
|
"gtSegments": {
|
"segments": [[(0.1, 0.1), (0.2, 0.2), (0.3, 0.3)]], # 一个分割掩码外轮廓
|
"labels": [2] # 分割标签
|
}
|
}
|
return {
|
"batch_inputs": x,
|
"ground_truth": gt
|
}
|
|
class ExampleModel(torch.nn.Module):
|
def __init__(self, mode="training"):
|
super(ExampleModel, self).__init__()
|
self.mode = mode
|
self.conv = torch.nn.Conv2d(3, 16, kernel_size=3, padding=1)
|
self.regressior = torch.nn.Linear(16, 4) # 假设输出3个类别
|
|
def forward(self, inputs):
|
batch_inputs = inputs["batch_inputs"]
|
ground_truth = inputs.get("ground_truth", None)
|
x = self.conv(batch_inputs) # 模拟前向传播
|
x = x.mean(dim=[2, 3], keepdim=True).squeeze(-1).permute(0,2,1) # 全局平均池化
|
boxes = self.regressior(x) # 回归输出
|
|
if self.training and self.mode == "training":
|
# 训练模式，计算损失
|
gt_boxes = torch.tensor([gt['gtBoxes']['boxes'] for gt in ground_truth])
|
box_loss = torch.nn.functional.mse_loss(boxes, gt_boxes.to(batch_inputs.device))
|
return {"losses": box_loss} # 返回损失
|
|
else:
|
pred_classes = torch.randint(0, 2, size=(boxes.size(0),)) # 模拟预测类别
|
pred_boxes = boxes
|
return {
|
"pred_output": {
|
"pred_classes": pred_classes.cpu(), # 预测类别
|
"pred_boxes": pred_boxes.cpu() # 预测框
|
},
|
"ground_truth": ground_truth
|
}
|
|
|
|
"""
|
- ExampleMetric
|
经过训练验证步骤后，将验证集中的所有预测结果和真值进行指标计算
|
|
-----call-----
|
** 输入：
|
验证集或测试集的所有预测结果和真值标签，若遵照上述ExampleModel的输出格式，则输入为一个字典列表
|
[{
|
"pred_output": {...},
|
"ground_truth": {...}
|
}, ...]
|
|
** 输出：
|
返回两个字典
|
metrics字典包含各类指标的计算结果，例如：
|
metrics ={
|
"检出率": 0.91, # 检测率
|
"过检率": 0.05, # 误检率
|
"像素级AUPRO": 0.8, # 像素级平均精度
|
"阈值": 0.5, # 阈值
|
"mAP@0.5": 0.87, # 平均精度均值
|
"mAP@0.5-0.95": 0.85, # 平均精度均值
|
}
|
predictions字典包含验证或测试集中所有图片的预测结果，例如：
|
predictions = {
|
4: { # 图片ID为4的预测结果
|
"imageId": 4, # 图片ID
|
"imagePath": "path/to/image_1.jpg", # 图片路径
|
"state": "漏检", # 检测状态["漏检", "误检", "正确", "未检"]
|
"iLevelPred": [1, 2], # 预测的图像级标签
|
"heatmapPath": "path/to/heatmap_1.jpg", # 热力图路径
|
"predBoxes": {
|
"boxes": [[0.1, 0.1, 0.5, 0.5]], # 预测框坐标
|
"classes": [1] # 框对应的标签
|
"scores": [0.9] # 框对应的置信度
|
}
|
"predSegments": {
|
"segments": [[[0.1, 0.1], [0.2, 0.2], [0.3, 0.3]]], # 预测分割掩码外轮廓
|
"classes": [2] # 分割标签
|
"scores": [0.8] # 分割置信度
|
}
|
}
|
}
|
"""
|
class ExampleMetric:
|
def __init__(self, mode="evaluation"):
|
self.mode = mode
|
|
def __call__(self, inputs):
|
metrics = {
|
"检出率": 0.91,
|
"过检率": 0.05,
|
"mAP@0.5": 0.87,
|
"像素级AUPRO": 0.8,
|
"阈值": 0.5,
|
}
|
|
if self.mode == "evaluation":
|
predictions = {}
|
return metrics, predictions
|
|
return metrics
|
|
def compare(self, metrics, best_metrics):
|
"""
|
比较当前评估指标与最佳指标
|
metrics: 当前评估指标
|
best_metrics: 最佳评估指标
|
返回True表示当前指标更好，False表示不如最佳指标
|
"""
|
if not best_metrics:
|
return True
|
# 假设我们只比较检出率
|
current_accuracy = metrics.get("检出率", 0)
|
best_accuracy = best_metrics.get("检出率", 0)
|
return current_accuracy > best_accuracy
|
|
|
@ModuleCreator.register("ExampleModel")
|
def create_module(config=None, mode="training"):
|
"""创建模拟模型"""
|
return ExampleModel()
|
|
@DataCreator.register("ExampleDataset")
|
def create_dataset(config=None, mode="training"):
|
"""创建模拟数据集"""
|
return ExampleDataset(config.train_data_path)
|
|
@OptimizerCreator.register("ExampleOptimizer")
|
def create_optimizer(model, config=None):
|
"""创建模拟优化器"""
|
return torch.optim.SGD(model.parameters(), lr=config.lr if config else 0.01)
|
|
@SchedulerCreator.register("ExampleScheduler")
|
def create_scheduler(optimizer, config=None):
|
"""创建模拟学习率调度器"""
|
scheduler_percent = config.scheduler_percent if config else (0.6, 0.9)
|
return torch.optim.lr_scheduler.MultiStepLR(optimizer, \
|
milestones=[int(p * config.epochs) for p in scheduler_percent], \
|
gamma=config.scheduler_gamma if config else 0.1,
|
last_epoch=config.start_epoch - 1 if config else -1)
|
|
@MetricorCreator.register("ExampleModel")
|
def create_metric(config=None, mode="evaluation"):
|
"""创建模拟评估器"""
|
return ExampleMetric()
|
