pytorch使用gpu运行

检测是否支持gpu

代码：

importtorch# 检查 CUDA 是否可用iftorch.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、计算 (自动在指定设备上运行)