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, # 自动调整批量大小以适应显存
|
||
)
|
||
|