当前位置：首页 > news >正文

Jupyter Notebook中使用GPU进行计算

news 来源：原创 2025/5/25 14:30:09

要在Jupyter Notebook中使用GPU进行计算，请按照以下步骤操作：

检查GPU可用性

在终端/Anaconda Prompt中运行以下命令以确认是否有可用的GPU：
```
nvidia-smi
```

安装所需的库

首先，在终端/Anaconda Prompt中，依次安装PyTorch及其相关的CPU和GPU支持包：

pip install torch torchvision torchaudio
# 如果需要多线程支持：
export OMP_NUM_THREADS=4
# 扩展CuDNN库的搜索路径（可选）：
LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-12.0/lib64

设置环境变量

在终端/Anaconda Prompt中，修改或新增.bashrc文件中的内容以配置以下环境变量（在每个新行后按Enter）：
```
LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-12.0/lib64
PATH=/usr/local/cuda-12.0/bin:$PATH
```
重新运行所有上述步骤可能需要多次登录别机或系统重启。

配置PyTorch使用GPU

在Jupyter Notebook中添加以下代码来配置 PyTorch 使用可用的 GPU：

import torch

if torch.cuda.is_available():
    print('Current device:', torch.cuda.current_device())
    print('Device name:', torch.cuda.get_device_name(0))
    # 设置多个可用GPU时，将索引从0开始设置
else:
    print('Using CPU!')

# 如果需要指定显卡：
torch.cuda.set_device(0)  # 替换为要使用的GPU索引

运行计算在GPU上

示例代码中创建两个张量 a 和 b 并执行GPU上的矩阵乘法：

a = torch.randn(2, 3).cuda()  # 将张量移动到GPU上
b = torch.randn(3, 4).cuda()
c = torch.mm(a, b)            # 使用PyTorch内置的矩阵乘法
d = a @ b                      # 添加Python 3.5+支持的短阵乘法

保存和重新载入修改
- 保存修改后的 Jupyter Notebook文件。
- 关闭前一个会话后，打开新会话并重新加载文件，否则新的代码更改可能不会生效。
验证计算是否在GPU上运行

运行以下命令查看 c 和 d 的设备：
```
print(f'张量 c位于{c.device}上')
print(f'张量 d位于{d.device}上')
```
检查是否有错误，比如 GPU可用性返回False或输出显示"cuda:0"。
处理常见问题
- 安装问题：确保 PyTorch 和_cuda版本与您的NVIDIA显卡匹配，并查阅[Pip Install guide](https://pytorch.org/docs/stable/get started.html# Install).
- 不支持的错误信息：如遇到无法导入模块，请检查PyTorch是否正确安装、路径是否设置正确或环境变量是否被正确识别。
- 无法使用GPU：确认命令是否执行后才指定 GPU（即修改代码后重新运行会话）。

扩展功能

如果需要，可以启用以下PyTorch选项以进一步优化 GPU 性能：

# 多GPU计算示例：
if torch.cuda.device_count() > 1:
    print("Let's use multiple GPUs!")
  
# 使用混合精度训练：
scaler = torch.cuda.amp.GradScaler()
  
# 启用自动混合精度（AMP）和执行优化选项：
with torch.amp.autocast():
    outputs = model(inputs)

optimizer.zero_grad()
loss.backward()
optimizer.step()