pytorch运行yolo5进行目标检测

1、安装anaconda

清华镜像源下载anaconda : mirrors.tuna.tsinghua.edu.cn/anaconda/ar…

配置环境变量

D:\dev\anaconda3-2024.10.1
D:\dev\anaconda3-2024.10.1\Scripts
D:\dev\anaconda3-2024.10.1\Library\bin
D:\dev\anaconda3-2024.10.1\Library\mingw-w64\bin

检查是否安装成功

conda --version

配置anaconda下载源 在 C:\Users\你的电脑用户名\ 下有一个 .condarc 文件，如果没有找到 .condarc ，输入以下命令

conda config --set show_channel_urls yes

打开 .condarc 复制下面内容保存

envs_dirs:
  - D:\dev\anaconda3-store\envs  #这是创建的虚拟环境路径，更改为自己的路径
pkgs_dirs:
  - D:\dev\anaconda3-store\pkgs  #这是anaconda的pkgs包路径，更改为自己的路径
channels:
  - defaults
custom_channels:
  conda-forge: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  msys2: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  bioconda: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  menpo: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  pytorch: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  simpleitk: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
default_channels:
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/msys2
show_channel_urls: True

检查配置

C:\Users\laibin>conda info

     active environment : None
       user config file : C:\Users\laibin\.condarc
 populated config files : D:\dev\anaconda3-2024.10.1\.condarc
                          C:\Users\laibin\.condarc
          conda version : 24.9.2
    conda-build version : 24.9.0
         python version : 3.12.7.final.0
                 solver : libmamba (default)
       virtual packages : __archspec=1=skylake
                          __conda=24.9.2=0
                          __cuda=12.5=0
                          __win=0=0
       base environment : D:\dev\anaconda3-2024.10.1  (read only)
      conda av data dir : D:\dev\anaconda3-2024.10.1\etc\conda
  conda av metadata url : None
           channel URLs : https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/win-64
                          https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/noarch
                          https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r/win-64
                          https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r/noarch
                          https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/msys2/win-64
                          https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/msys2/noarch
                          https://repo.anaconda.com/pkgs/main/win-64
                          https://repo.anaconda.com/pkgs/main/noarch
                          https://repo.anaconda.com/pkgs/r/win-64
                          https://repo.anaconda.com/pkgs/r/noarch
                          https://repo.anaconda.com/pkgs/msys2/win-64
                          https://repo.anaconda.com/pkgs/msys2/noarch
          package cache : D:\dev\anaconda3-store\pkgs
       envs directories : D:\dev\anaconda3-store\envs
                          C:\Users\laibin\.conda\envs
                          D:\dev\anaconda3-2024.10.1\envs
                          C:\Users\laibin\AppData\Local\conda\conda\envs
               platform : win-64
             user-agent : conda/24.9.2 requests/2.32.3 CPython/3.12.7 Windows/11 Windows/10.0.22631 solver/libmamba conda-libmamba-solver/24.9.0 libmambapy/1.5.8 aau/0.4.4 c/. s/.
          administrator : False
             netrc file : None
           offline mode : False

2、安装pytorch

创建一个虚拟环境

conda create --name pytorch-env python=3.12
conda activate pytorch-env

进入 pytorch 官网 : pytorch.ac.cn/ ，选择自己下载的版本

在 pytorch-env 环境下载 pytorch

conda install pytorch torchvision torchaudio pytorch-cuda=12.4 -c pytorch -c nvidia

下载完毕，验证是否可用，切换到 pytorch-env 环境中

import torch

if __name__ == '__main__':
    print(torch.__version__)
    print(f"CUDA available: {torch.cuda.is_available()}")
    print(f"Current device: {torch.cuda.current_device()}")
    print(f"device name: {torch.cuda.get_device_name()}")

# --------输出
#2.5.1
#CUDA available: True
#Current device: 0
#device name: NVIDIA GeForce RTX 4070 Laptop GPU

3、运行yolo5

下载 yolo5 源码 : github.com/ultralytics…

修改 models/experimental.py ，使用GPU

model = Ensemble()
    for w in weights if isinstance(weights, list) else [weights]:
        # ckpt = torch.load(attempt_download(w), map_location="cpu")  # load
        ckpt = torch.load(attempt_download(w), map_location="cuda:0") # 使用 GPU
        ckpt = (ckpt.get("ema") or ckpt["model"]).to(device).float()  # FP32 model

        # Model compatibility updates
        if not hasattr(ckpt, "stride"):
            ckpt.stride = torch.tensor([32.0])
        if hasattr(ckpt, "names") and isinstance(ckpt.names, (list, tuple)):
            ckpt.names = dict(enumerate(ckpt.names))  # convert to dict

        model.append(ckpt.fuse().eval() if fuse and hasattr(ckpt, "fuse") else ckpt.eval())  # model in eval mode

运行 detect.py ，结果保存在 runs/detect/ 目录下