学习记录 DAY7 正确使用CUDA跑yolo的模型

上述在安装CUDA的时候，其实我们的操作都是正确的，而且在默认勾选的时候已经勾选了注入对应的环境变量到本系统的PATH中，

这里是可以看到已经有正确路径的。

但是我们在使用检测python是否能用CUDA时，输入以下代码：

import torch

# 检查 CUDA 是否可用
print(torch.cuda.is_available())

会发现输出的结果是

2.6.0+cpu

也就是说，这里是cpu版本的pytorch，因此并不会用到CUDA。

所以我们需要重新卸载pytorch，并且安装支持CUDA版本的pytorch。

首先卸载当前pytorch

pip uninstall torch torchvision torchaudio

然后，我们需要清除掉对应的pip缓存

pip cache purge

这里要注意了，当我们再次安装Pytorch的时候，一定要注意，pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126 -i https://pypi.tuna.tsinghua.edu.cn/simple这个路径下载的，可能会因为走的路径的原因，一直给你下的是CPU的版本！也因此无法使用CUDA！！！！

然后我问了问学长+种种各种原因，我决定用12.4版本的CUDA，然后精确拿下路径下载。

所以，我们需要用这个命令来精确地下载：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 -i https://pypi.tuna.tsinghua.edu.cn/simple

这个是CUDA12.4

如果当你安装的时候看到安装的东西大小是2.5G（也有可能是2.4G，我可能记错，反正差不多这个大小）

那你基本上就会安装对了。此外，我的python解释器用的是3.9版本，也是大伙公认的稳定的版本。

当我们在python中test一个py文件，输入以下代码：

import torch

# 检查 CUDA 是否可用
print(torch.cuda.is_available())

if torch.cuda.is_available():
    # 获取可用的 CUDA 设备数量
    print("可用的 CUDA 设备数量:", torch.cuda.device_count())
    # 获取当前使用的 CUDA 设备索引
    print("当前 CUDA 设备索引:", torch.cuda.current_device())
    # 获取指定索引 CUDA 设备的名称
    print("当前 CUDA 设备名称:", torch.cuda.get_device_name(0))

    # 创建一个张量并将其移动到 CUDA 设备上
    x = torch.tensor([1.0, 2.0]).cuda()
    y = torch.tensor([3.0, 4.0]).cuda()
    z = x + y
    print("在 CUDA 上计算的结果:", z)
else:
    print("CUDA 不可用")

如果提示能用，那就大功告成啦。

然后我的数据集使用的是FER-2013的人脸识别表情的图库。

没有labels，然后在cursor的帮助下弄了个打标签的py文件，先运行这个。

import os
import cv2
import numpy as np


def create_yolo_labels(dataset_path, class_id):
    """为每个图像创建YOLO格式的标签文件"""
    images = os.listdir(dataset_path)

    for img_name in images:
        if not img_name.endswith(('.jpg', '.png', '.jpeg')):
            continue

        # 读取图像获取尺寸
        img_path = os.path.join(dataset_path, img_name)
        img = cv2.imread(img_path)
        if img is None:
            continue

        h, w = img.shape[:2]

        # 创建标签文件
        label_path = os.path.join(dataset_path, img_name.rsplit('.', 1)[0] + '.txt')

        # YOLO格式：<class> <x_center> <y_center> <width> <height>
        # 假设人脸占据整个图像的中心区域
        with open(label_path, 'w') as f:
            f.write(f"{class_id} 0.5 0.5 0.8 0.8\n")


def process_dataset(base_path):
    """处理整个数据集"""
    emotions = {
        'angry': 0,
        'disgust': 1,
        'fear': 2,
        'happy': 3,
        'sad': 4,
        'surprise': 5,
        'neutral': 6
    }

    # 处理训练集和测试集
    for split in ['train', 'test']:
        split_path = os.path.join(base_path, split)
        for emotion, class_id in emotions.items():
            emotion_path = os.path.join(split_path, emotion)
            if os.path.exists(emotion_path):
                print(f"Processing {split}/{emotion}")
                create_yolo_labels(emotion_path, class_id)


if __name__ == '__main__':
    # 设置数据集路径
    dataset_path = "D:/DeskTop/ultralytics-main/datasets/archive"
    process_dataset(dataset_path)

然后，我们就可以在train.py中跑自己指定的模型来训练了。

from ultralytics import YOLO
import torch
from multiprocessing import freeze_support

def main():
    # 检查CUDA是否可用
    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
    print(f"Using device: {device}")

    # 加载预训练模型
    model = YOLO('yolov8n.pt')  # 使用较小的n模型以适应4G显存

    # 训练参数设置
    results = model.train(
        data='ultralytics/cfg/datasets/emotion_dataset.yaml',  # 数据集配置文件的相对路径
        epochs=100,                   # 训练轮数
        imgsz=64,                    # 必须是32的倍数
        batch=16,                    # 较小的batch size以适应显存
        patience=20,                 # 早停参数
        save_period=10,             # 每10轮保存一次
        device=device,              # 自动选择设备
        workers=0,                  # 减少worker数量，避免多进程问题
        cache=False,                # 关闭缓存以节省内存
        project='runs/emotion',     # 项目保存路径
        name='exp1',               # 实验名称
        pretrained=True,           # 使用预训练权重
        optimizer='Adam',          # 优化器选择
        lr0=0.001,                # 初始学习率
        weight_decay=0.0005,      # 权重衰减
        warmup_epochs=3,          # 预热轮数
        close_mosaic=10,          # 最后10轮关闭马赛克增强
        resume=False,             # 是否继续训练
        single_cls=True           # 单类别模式
    )

if __name__ == '__main__':
    freeze_support()
    main()

我的显卡是3050ti,4G显存，比较拉胯，所以跑的强度不是很高，但是也依然要一轮约4分钟，总共跑100轮大约需要6个多小时，目前仍在跑中。

之后关于训练结果使用，UI建立，就得等到明天了，我得去看看深度学习的MOOC了