26 lines
852 B
Python
26 lines
852 B
Python
|
from ultralytics import YOLO
|
|||
|
|
import torch
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 检查CUDA是否可用
|
|||
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|||
|
|
print(f"Using device: {device}")
|
|||
|
|
|
|||
|
|
# 加载模型
|
|||
|
|
model = YOLO("runs/detect/train6/weights/last.pt").to(device)
|
|||
|
|
|
|||
|
|
# 设置新的分辨率
|
|||
|
|
imgsz = 1024 # 这里将图像尺寸调整为 1280x1280（你可以根据显存调整尺寸）
|
|||
|
|
|
|||
|
|
# 训练模型，传入增强参数
|
|||
|
|
model.train(
|
|||
|
|
data="dataset/dataset.yaml", # 你的数据集配置文件
|
|||
|
|
epochs=1000, # 训练轮次
|
|||
|
|
imgsz=imgsz, # 使用更高的分辨率
|
|||
|
|
device=[1], # 使用第一块 GPU（如果有多个 GPU，可以调整）
|
|||
|
|
hsv_v=0.3, # 修改图像亮度的一部分，帮助模型在不同光照条件下表现良好
|
|||
|
|
cos_lr=True, # 启用余弦学习率调度
|
|||
|
|
batch = -1, # 自动调整批量大小以适应显存
|
|||
|
|
)
|
|||
|
|
|