sun15684712291/aicheckv2
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

使用cuda版torch进行训练和推理

Browse Source
This commit is contained in:
sunyg
2025-03-17 15:45:46 +08:00
parent ff7b6fda0e
commit ebc4e9df4f
5 changed files with 37 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
requirements_gpu.txt
View File

@ -24,7 +24,7 @@ python-socketio == 5.12.1
# BASE ------------------------------------------------------------------------
gitpython>=3.1.30
matplotlib>=3.3
numpy==2.0.2
numpy==1.24.0 # cuda版本的numpy需要降一下版本
opencv-python>=4.1.1
pillow>=10.3.0
psutil # system resources
@ -32,12 +32,11 @@ PyYAML>=5.3.1
requests>=2.32.2
scipy==1.13.1
thop>=0.1.1 # FLOPs computation
torch --extra-index-url https://download.pytorch.org/whl/cu121
torchvision --extra-index-url https://download.pytorch.org/whl/cu121
torchaudio --extra-index-url https://download.pytorch.org/whl/cu121
torch==2.3.1+cu121 --extra-index-url https://mirrors.aliyun.com/pytorch-wheels/cu121/
torchvision==0.18.1+cu121 --extra-index-url https://mirrors.aliyun.com/pytorch-wheels/cu121/
tqdm>=4.66.3
ultralytics>=8.2.34 # https://ultralytics.com
ultralytics==8.3.75 # https://ultralytics.com
# Plotting --------------------------------------------------------------------
pandas==2.2.3
seaborn>=0.11.0
seaborn==0.11.0 # 对应的这个依赖也需要降