130 lines
6.5 KiB
Python
130 lines
6.5 KiB
Python
import yaml, sys
|
||
import torch
|
||
from pathlib import Path
|
||
|
||
# --- 1. 核心目录设置 (Core Directories) ---
|
||
HARDISK_DIR = Path.home() / "Desktop" / "Seg" / "Hardisk" # 硬盘根目录
|
||
BASE_DIR = Path(__file__).parent.parent # 当前脚本所在目录上级
|
||
DATASET_YAML_PATH = Path(__file__).parent / "dataset.yaml" # 数据集的 YAML 配置文件路径
|
||
OUTPUTS_DIR = BASE_DIR / "DataSet_Public_outputs" # 输出结果目录 # 最终为：OUTPUTS_DIR / dataset_name
|
||
|
||
TEST_IMAGE_DIR = None # 初始化 TEST_IMAGE_DIR # 测试文件地址 dataset.path / TODO "val" / "test" TODO
|
||
# try:
|
||
# 3. 读取并解析 YAML 文件
|
||
with open(DATASET_YAML_PATH, 'r', encoding='utf-8') as f:
|
||
yaml_data = yaml.safe_load(f)
|
||
# 4. 从解析后的数据中获取 'path' 的值
|
||
relative_path_from_yaml = yaml_data.get('path')
|
||
test_path_from_yaml = yaml_data.get('test')
|
||
dataset_name = Path(relative_path_from_yaml).name
|
||
if relative_path_from_yaml:
|
||
# 5. 【核心步骤】构建绝对路径
|
||
# relative_path_from_yaml 是相对于 .yaml 文件本身的路径。
|
||
# 所以，我们需要获取 .yaml 文件所在的目录，然后与这个相对路径拼接。
|
||
yaml_file_directory = DATASET_YAML_PATH.parent
|
||
TEST_IMAGE_DIR = (DATASET_YAML_PATH.parent / relative_path_from_yaml / test_path_from_yaml).resolve() # 这里val 或 test在dataset.yaml中定义
|
||
# 使用 .exists() 方法来检查路径是否存在
|
||
if not TEST_IMAGE_DIR.exists():
|
||
# 如果路径不存在，则执行这里的代码
|
||
print(f"警告: 测试图片目录不存在: {TEST_IMAGE_DIR}")
|
||
sys.exit(1)
|
||
# 6. 获取OUTPUTS_DIR
|
||
if dataset_name and dataset_name not in {'.', '..'}:
|
||
dataset_name_ = dataset_name + "-Yolo"
|
||
OUTPUTS_DIR = OUTPUTS_DIR / dataset_name_
|
||
print(f"设定输出路径为: '{str(OUTPUTS_DIR)}'")
|
||
else:
|
||
print(f"警告: 提取的dataset_name: '{dataset_name}' 无效（为空、'.' 或 '..'）。")
|
||
sys.exit(1)
|
||
else:
|
||
print(f"警告: 在 '{DATASET_YAML_PATH}' 文件中没有找到 'path' 键。")
|
||
sys.exit(1)
|
||
# 4. 从解析后的数据中获取 'path' 的值
|
||
relative_path_from_yaml = yaml_data.get('path')
|
||
# except FileNotFoundError:
|
||
# print(f"错误: YAML 配置文件未找到: '{DATASET_YAML_PATH}'")
|
||
# except Exception as e:
|
||
# print(f"读取或解析 YAML 文件时发生错误: {e}")
|
||
|
||
|
||
# ==============================================================================
|
||
# --- 新增：模型选择 (只需修改这里) ---
|
||
# 你想要使用的模型名称，从下面的 MODEL_CONFIGS 字典中选择一个键。
|
||
SELECTED_MODEL = 'YOLOv9e-seg'
|
||
# ==============================================================================
|
||
|
||
|
||
# --- 修改：模型、图像大小和批处理大小的集成配置 ---
|
||
# 将所有模型的配置（权重、图像大小、批处理大小）集中管理
|
||
MODEL_CONFIGS = {
|
||
# YOLOv8
|
||
'YOLOv8n-seg': {'weights': 'yolov8n-seg.pt', 'image_size': 640, 'batch_size': 16},
|
||
'YOLOv8s-seg': {'weights': 'yolov8s-seg.pt', 'image_size': 640, 'batch_size': 16},
|
||
'YOLOv8m-seg': {'weights': 'yolov8m-seg.pt', 'image_size': 640, 'batch_size': 16},
|
||
'YOLOv8l-seg': {'weights': 'yolov8l-seg.pt', 'image_size': 640, 'batch_size': 16}, # 640，16，12.5GB
|
||
'YOLOv8x-seg': {'weights': 'yolov8x-seg.pt', 'image_size': 640, 'batch_size': 16}, # 640，16，15.5GB # 示例：X模型使用1280分辨率
|
||
# YOLOv9
|
||
'YOLOv9c-seg': {'weights': 'yolov9c-seg.pt', 'image_size': 640, 'batch_size': 16}, # 640，16，13GB
|
||
'YOLOv9e-seg': {'weights': 'yolov9e-seg.pt', 'image_size': 640, 'batch_size': 8}, # 640，16，内存超了 # 示例：E模型（最大）使用1280分辨率和更小的batch
|
||
# YOLOv11 (假设)
|
||
'YOLO11n-seg': {'weights': 'yolo11n-seg.pt', 'image_size': 640, 'batch_size': 16},
|
||
'YOLO11s-seg': {'weights': 'yolo11s-seg.pt', 'image_size': 640, 'batch_size': 16},
|
||
'YOLO11m-seg': {'weights': 'yolo11m-seg.pt', 'image_size': 640, 'batch_size': 16},
|
||
'YOLO11l-seg': {'weights': 'yolo11l-seg.pt', 'image_size': 640, 'batch_size': 16}, # 640，16，12.5GB
|
||
'YOLO11x-seg': {'weights': 'yolo11x-seg.pt', 'image_size': 640, 'batch_size': 16}, # 640，16，19.5GB
|
||
# YOLOv12 (假设)
|
||
'YOLO12-seg': {'weights': str(Path(__file__).parent / 'yolo12-seg.yaml'), 'image_size': 640, 'batch_size': 16}, # 640，16，3GB
|
||
}
|
||
|
||
# --- 新增：根据选择自动加载配置 ---
|
||
# 检查所选模型是否存在于配置字典中
|
||
if SELECTED_MODEL in MODEL_CONFIGS:
|
||
config = MODEL_CONFIGS[SELECTED_MODEL]
|
||
MODEL_WEIGHTS = config['weights']
|
||
IMAGE_SIZE = config['image_size']
|
||
BATCH_SIZE = config['batch_size']
|
||
print(f"--- 模型配置加载成功 ---")
|
||
print(f"模型名称: {SELECTED_MODEL}")
|
||
print(f"权重文件: {MODEL_WEIGHTS}")
|
||
print(f"图像大小: {IMAGE_SIZE}")
|
||
print(f"批处理大小: {BATCH_SIZE}")
|
||
print(f"------------------------")
|
||
else:
|
||
print(f"错误: 所选模型 '{SELECTED_MODEL}' 不在配置列表中。")
|
||
print("可用模型包括: ", list(MODEL_CONFIGS.keys()))
|
||
sys.exit(1)
|
||
|
||
|
||
# --- 4. 训练超参数 (Training Hyperparameters) ---
|
||
EPOCHS = 300 # 训练轮次
|
||
PATIENCE = 100 # 提前停止训练的轮数
|
||
SAVE_PERIOD = 10 # 每隔多少轮保存一次模型
|
||
# BATCH_SIZE 已在上面根据模型自动设置
|
||
LEARNING_RATE = 0.01 # 初始学习率
|
||
OPTIMIZER = 'Adam' # 优化器 (TODO 'SGD', 'Adam', 'AdamW', 'auto' TODO)
|
||
WORKERS = 4 # 数据加载的工作线程数
|
||
# --- 动态设备选择 (Dynamic Device Selection) ---
|
||
def get_auto_device():
|
||
"""自动检测并返回最合适的设备"""
|
||
if torch.cuda.is_available():
|
||
gpu_count = torch.cuda.device_count()
|
||
if gpu_count > 1:
|
||
# 如果有多个GPU，使用所有GPU
|
||
device_str = ",".join(str(i) for i in range(gpu_count))
|
||
print(f"检测到 {gpu_count} 个可用的GPU。将使用所有GPU: {device_str}")
|
||
return device_str
|
||
else:
|
||
# 如果只有1个GPU
|
||
print("检测到 1 个可用的GPU。将使用 GPU: 0")
|
||
return '0'
|
||
else:
|
||
# 如果没有可用的GPU，使用CPU
|
||
print("未检测到可用的GPU。将使用CPU。")
|
||
return 'cpu'
|
||
DEVICE = get_auto_device() # 调用函数来设置设备
|
||
|
||
# --- 5. 预测设置 (Prediction Settings) ---
|
||
|
||
SHOW_LABELS = True # 是否在预测结果上显示类别标签
|
||
SHOW_CONF = True # 是否在预测结果上显示置信度和边界框
|
||
SAVE_PREDICTIONS = True # 是否保存预测结果图像 |