PyTorch torch.cuda 函数
Pytorch torch 参考手册torch.cuda 是 PyTorch 中用于 CUDA 管理的模块。提供了 CUDA 设备管理、内存管理、同步等功能。
函数定义
torch.cuda
torch.cuda 模块包含以下常用函数:
torch.cuda.is_available()- 检查 CUDA 是否可用torch.cuda.device_count()- 获取 CUDA 设备数量torch.cuda.current_device()- 获取当前设备torch.cuda.synchronize()- 同步 CUDA 操作torch.cuda.empty_cache()- 清空缓存
使用示例
实例
import torch
# 检查 CUDA 是否可用
print(f"CUDA available: {torch.cuda.is_available()}")
# 获取 CUDA 设备数量
if torch.cuda.is_available():
print(f"CUDA device count: {torch.cuda.device_count()}")
print(f"CUDA device name: {torch.cuda.get_device_name(0)}")
# 检查 CUDA 是否可用
print(f"CUDA available: {torch.cuda.is_available()}")
# 获取 CUDA 设备数量
if torch.cuda.is_available():
print(f"CUDA device count: {torch.cuda.device_count()}")
print(f"CUDA device name: {torch.cuda.get_device_name(0)}")