yolov8 训练后写脚本

以下是使用YOLOv8训练后编写脚本的示例[^1][^2]：

import torch
from models import Darknet
from utils.datasets import LoadImagesAndLabels
from utils.general import check_img_size, non_max_suppression, scale_coords
from utils.torch_utils import select_device, time_synchronized

# 设置设备
device = select_device('0')

# 加载模型
model = Darknet('cfg/yolov3.cfg', img_size=416).to(device)
model.load_state_dict(torch.load('weights/yolov3.pt', map_location=device)['model'])
model.eval()

# 加载数据集
dataset = LoadImagesAndLabels('data/train.txt', img_size=416, batch_size=8)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=8, shuffle=True, num_workers=4)

# 训练脚本
for epoch in range(100):
    for batch_i, (imgs, targets, paths, shapes) in enumerate(dataloader):
        imgs = imgs.to(device)
        targets = targets.to(device)

        # 前向传播
        with torch.no_grad():
            outputs = model(imgs)

        # 后处理
        outputs = non_max_suppression(outputs, conf_thres=0.5, iou_thres=0.5)

        # 打印结果
        for output in outputs:
            if output is not None:
                output[:, :4] = scale_coords(imgs.shape[2:], output[:, :4], shapes)
                for x1, y1, x2, y2, conf, cls_conf, cls in output:
                    print(f'预测结果：{cls.item()}')

        # 反向传播
        loss = model(imgs, targets)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        # 打印损失
        print(f'Epoch {epoch}, Batch {batch_i}, Loss: {loss.item()}')

请注意，上述代码仅为示例，具体的脚本可能需要根据您的数据集和模型进行适当的修改。