### PyTorch 中使用 GPU 加速模型训练或推理 在 Python 的 PyTorch 框架中，可以通过以下方式利用 GPU 来加速模型的训练或推理过程。 #### 判断 GPU 是否可用 首先需要确认当前环境中是否存在可使用的 GPU。这一步非常重要，因为如果设备不支持 CUDA，则无法启用 GPU 计算。可以使用 `torch.cuda.is_available()` 函数来检测是否有可用的 GPU[^2]。 ```python import torch if torch.cuda.is_available(): print("CUDA is available! Using GPU.") else: print("No CUDA detected. Falling back to CPU.") ``` #### 设置设备对象 为了使代码能够在不同硬件环境下无缝切换（即无论是在 CPU 或 GPU 上运行），通常会定义一个通用的设备对象 `device`。该对象可以根据实际条件自动选择合适的计算资源： ```python device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") ``` 此部分逻辑已在多个引用中提及[^3]。 #### 数据转移至 GPU 对于张量和模型而言，在执行任何操作前都需要将其显式地移动到目标设备上。这是通过调用 `.to(device)` 方法实现的。相比于旧版 API 提供的 `.cuda()` 接口，推荐优先采用更灵活且兼容性更好的 `.to()` 方式[^4]。 以下是具体示例代码片段展示如何完成这一迁移过程： ```python # 假设我们有一个简单的线性层作为模型实例化 model = torch.nn.Linear(10, 5) # 将模型放置到选定设备之上 model.to(device) # 创建随机输入张量并同样迁移到对应位置 input_tensor = torch.randn(8, 10).to(device) output_tensor = model(input_tensor) print(output_tensor.shape) ``` 注意这里不仅涉及到了模型本身还需要处理所有关联的数据结构如批次样本等等均需同步调整其存储地址以便后续正常交互运算[^3]. #### 并行多GPU配置 当存在超过一块独立图形处理器可供调度时，还可以进一步优化性能表现借助分布式技术手段充分利用这些额外能力提升整体效率水平。其中最常见的一种做法便是运用内置模块 `nn.DataParallel`, 它允许轻松扩展单一网络架构跨数台机器协同作业而无需修改太多原有设计思路[^5]: ```python os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2" if torch.cuda.device_count() > 1: print(f"{torch.cuda.device_count()} GPUs are used!") model = nn.DataParallel(model) model.to(device) ``` 以上脚本展示了怎样重新映射物理编号体系从而避开默认行为带来的干扰问题同时也体现了动态分配策略的优势所在. --- ### 总结 综上所述，要成功部署基于 PyTorch 构建的应用程序使之受益于现代高性能计算单元所提供的强大动力源——图形处理器(GPU)，主要分为以下几个方面的工作要点：首先是验证基础环境准备情况；其次是合理规划资源配置方案最后才是编写清晰易懂易于维护升级版本控制友好的高质量源码文件[^1]-[^5].