手写最基础的yolov5训练代码

以下是一个基础的YOLOv5训练代码示例：

import torch
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
from models.experimental import attempt_load
from utils.general import check_dataset, check_file, check_img_size

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 数据集路径
data_path = "path/to/dataset"

# 检查数据集是否合法
check_dataset(data_path)

# 加载数据集
dataset = datasets.ImageFolder(data_path, transform=ToTensor())
dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

# 加载模型配置文件和权重文件
model_cfg = "path/to/model.yaml"
model_weights = "path/to/weights.pt"

# 检查配置文件和权重文件是否存在
check_file(model_cfg)
check_file(model_weights)

# 加载模型
model = attempt_load(model_cfg, model_weights, device)

# 设置图片尺寸
img_size = check_img_size(640)

# 设置优化器
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# 设置损失函数
criterion = torch.nn.CrossEntropyLoss()

# 训练循环
num_epochs = 10

for epoch in range(num_epochs):
    for batch, (images, labels) in enumerate(dataloader):
        # 将数据移动到设备
        images = images.to(device)
        labels = labels.to(device)
        
        # 前向传播
        outputs = model(images)
        
        # 计算损失
        loss = criterion(outputs, labels)
        
        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # 输出训练信息
        if batch % 10 == 0:
            print(f"Epoch {epoch+1}/{num_epochs}, Batch {batch+1}/{len(dataloader)}, Loss: {loss.item()}")

在上述代码中，我们首先设置了设备（CPU或GPU），然后加载数据集和模型。接下来，我们设置优化器和损失函数，并进行训练循环。在每个批次中，我们将数据移动到设备上，执行前向传播、计算损失、反向传播和优化，并输出训练信息。

请注意，这只是一个基础的YOLOv5训练代码示例，实际使用时可能需要根据具体的数据集、模型和训练需求进行相应的调整和修改。