pytorch框架训练模型GPU加速：手把手教程

pytorch框架训练模型GPU加速：手把手教程

前言：
训练如果使用cpu太慢了。屏幕前的你，是否烦恼于CUDA、cuDNN、Pytorch版本不一致的问题？是否烦恼于虚拟环境配置Pytorch和CUDA的问题？
OK，just follow me.
1： 你先确定好你要搞的代码使用的是什么版本的python环境、pytorch版本【一般readme.md or requirements文件中都会有介绍】
2： 检查电脑的CUDA驱动【允许你安装什么版本】
3： 把你需要的torch版本+你的CUDA版本进行对应确定
4： 安装CUDA【配置path（图形界面安装的一般会自动配置）】
5： 安装cuDNN【注意安装CUDA对应的版本】
6： 创建你要搞的代码的虚拟环境【conda】
7： 在虚拟环境中安装torch（比如CUDA12.1：https://download.pytorch.org/whl/cu121）
8： 在训练代码中添加device=0调用GPU

1：拿到代码，确定环境：比如YOLOv12

2：检查CUDA驱动【一般向下兼容】

3：对应版本
https://pytorch.org/

https://developer.nvidia.com/cuda-toolkit-archive

安装CUDA很简单，下一步下一步（略过）
安完之后：nvcc --version

然后安装cuDNN：https://developer.nvidia.com/cudnn-archive

然后conda创建虚拟环境，激活虚拟环境
pip安装torch：https://download.pytorch.org/whl/cu121
会有点慢，不如直接这个网站下载，然后再pip本地安装：


安完之后，我一运行，提示我需要torchvision包，那就再安装：（修正：不要直接pip install torchvision，改为和torch统一版本，在网站上下载本地安装）

安了之后，小心它把你的可GPU加速的torch给卸了，偷偷安了个CPU加速的torch，这时候，再执行上面的安装你下载好的本地的torch即可。
最后：

import torch

# 检查 CUDA 是否可用
print(f"CUDA available: {torch.cuda.is_available()}")

print(torch.__version__)

print(f"CUDA version: {torch.version.cuda}")

# 检查 cuDNN 是否可用
if torch.backends.cudnn.is_available():
    print("cuDNN is available.")
    print(f"cuDNN version: {torch.backends.cudnn.version()}")
else:
    print("cuDNN is not available.")

# 尝试在 GPU 上创建一个张量
if torch.cuda.is_available():
    device = torch.device("cuda")
    x = torch.tensor([1.0], device=device)
    print(f"Tensor on GPU: {x}")

输出：

OK了，接下来你就随便玩弄GPU吧【加上device=0即可】：

from ultralytics import YOLO
import torch

def main():
    model = YOLO(r'ultralytics\cfg\models\12\yolo12-obb.yaml').load(
        r'model\yolo11n-obb.pt')  # build from YAML and transfer weights
    model.train(data=r'ultralytics\cfg\datasets\DOTAv1.5.yaml', epochs=100, imgsz=640, batch=4, workers=0, device=0)


if __name__ == '__main__':
    # print(torch.cuda.is_available())  # 如果cuda可用，则返回True
    main()