PyTorch CUDA 使用教程：加速模型训练与推理

在 PyTorch 中使用 CUDA 可以实现在 GPU 上加速模型训练和推理过程。下面是一个示例，展示了如何在 PyTorch 中使用 CUDA：

首先，确保您的 PyTorch 版本和 CUDA 版本兼容。您可以在 PyTorch 官方网站上查看兼容性信息。

1. 导入 PyTorch 和相关模块：

import torch
import torch.nn as nn

2. 检查 CUDA 是否可用，并设置设备：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

3. 定义模型并将其移动到所选设备：

model = YourModel().to(device)

4. 在训练和推理过程中，将输入数据和目标数据移动到设备上：

inputs = inputs.to(device)
targets = targets.to(device)

5. 在每个迭代中，将输入数据传递给模型时，确保将模型设置为训练模式或评估模式：

# 训练模式
model.train()

# 评估模式
model.eval()

6. 在模型的前向传播中，确保使用 CUDA 加速计算：

outputs = model(inputs)

7. 如果有需要，您还可以手动将特定的张量移动到设备上：

tensor = tensor.to(device)

这是一个简单的示例，展示了如何在 PyTorch 中使用 CUDA。根据您的具体需求，您可能需要进行更复杂的操作，如使用多个 GPU 或使用 CUDA 相关函数。请参考 PyTorch 的官方文档以获取更详细的信息和示例。