yolov8训练代码python

以下是YOLOv3的训练代码示例，YOLOv8的代码与之类似：

import argparse
import os
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms
from dataset import CustomDataset
from model import YOLOv3

def train(model, dataloader, optimizer, criterion, device):
    model.train()
    running_loss = 0.0
    for images, targets in dataloader:
        images = images.to(device)
        targets = targets.to(device)
        
        optimizer.zero_grad()
        
        outputs = model(images)
        loss = criterion(outputs, targets)
        
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
    
    epoch_loss = running_loss / len(dataloader)
    return epoch_loss

def main(args):
    # 设置GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 创建数据集和数据加载器
    train_transforms = transforms.Compose([
        transforms.Resize((416, 416)),
        transforms.ToTensor()
    ])
    dataset = CustomDataset(args.data_dir, args.annotation_dir, transform=train_transforms)
    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)

    # 创建模型
    model = YOLOv3(num_classes=args.num_classes).to(device)

    # 定义损失函数和优化器
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)

    # 开始训练
    for epoch in range(args.num_epochs):
        epoch_loss = train(model, dataloader, optimizer, criterion, device)
        print(f"Epoch {epoch+1}/{args.num_epochs}, Loss: {epoch_loss}")

    # 保存模型
    torch.save(model.state_dict(), args.save_path)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_dir", type=str, default="data/")
    parser.add_argument("--annotation_dir", type=str, default="annotations/")
    parser.add_argument("--batch_size", type=int, default=8)
    parser.add_argument("--num_epochs", type=int, default=10)
    parser.add_argument("--num_classes", type=int, default=80)
    parser.add_argument("--learning_rate", type=float, default=0.001)
    parser.add_argument("--save_path", type=str, default="model.pth")
    args = parser.parse_args()

    main(args)

请注意，这只是一个示例代码，需要根据你的数据集和网络架构进行适当的修改。确保你已经准备好了训练图像和相应的注释文件，并将它们放在正确的目录中。你还需要定义自己的模型和损失函数，根据你的任务进行必要的更改。