pytorch使用gpu运行 todo抽空迁移并废弃掉
pytorch使用gpu运行
·
文章目录
pytorch默认是使用cpu的,要使用gpu需要设置下。
检测是否支持gpu
代码:
import torch
# 检查 CUDA 是否可用
if torch.cuda.is_available():
print(f"GPU 可用!显卡型号: {torch.cuda.get_device_name(0)}")
print(f"显存总量: {torch.cuda.get_device_properties(0).total_memory / 1024**2:.2f} MB")
else:
print("GPU 不可用,将回退到 CPU。请检查:\n1. 是否安装了 NVIDIA 驱动\n2. 是否安装了 cuda 版本的 pytorch (pip install torch --index-url https://download.pytorch.org/whl/cu118)")
gpu和cpu切换使用
1、检查 torch.cuda.is_available()
2、定义 device = torch.device(“cuda” if … else “cpu”)
3、移动模型 model.to(device)
4、移动数据 data = data.to(device) (在每一个 training step 循环里都要做)
5、计算 (自动在指定设备上运行)
更多推荐
所有评论(0)