如何在Jupyter Notebook中指定使用特定GPU设备

首先直接回答export CUDA_VISIBLE_DEVICES对Jupyter是否生效的问题：
只有在启动Jupyter进程前设置才生效，启动后再在其他终端设置无效。环境变量是进程级的，Jupyter启动时会继承当前终端的环境，之后你在别的终端改CUDA_VISIBLE_DEVICES，根本影响不到已经跑起来的Jupyter进程，自然不生效。

下面是几种亲测可用的指定GPU方法，按可靠程度排序：

1. 启动前全局指定（最稳妥，零代码侵入） #

第一步先敲nvidia-smi看GPU占用，找显存、算力都空闲的卡，记好物理编号，比如3号卡没人用。
在你准备启动Jupyter的终端里，先执行：

export CUDA_VISIBLE_DEVICES=3
# 要同时用多张空闲卡就用逗号分隔编号，比如export CUDA_VISIBLE_DEVICES=3,4

确认变量生效可以敲echo $CUDA_VISIBLE_DEVICES看输出是不是你设的编号，确认没问题之后，直接在同一个终端敲你平时用的启动命令就行，不管是jupyter notebook还是jupyter lab都可以。

注意：别开两个终端操作，一个终端设变量、另一个终端启Jupyter是没用的，必须在同一个会话里先设变量再启动。
这种方法启动的Jupyter，所有新开的Notebook内核默认都只能看到你指定的GPU，完全不会访问其他卡上别人跑的任务，适合长期用固定卡跑任务的场景。

2. Notebook内指定（不用重启Jupyter服务） #

要是你已经把Jupyter启动了，不想重启服务中断手头的东西，就打开你要跑的.ipynb文件，在所有导入深度学习框架的代码之前（也就是import torch、import tensorflow这类语句前面），新加一个代码块写：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
# 多卡就写"3,4"

⚠️ 位置一定要放在框架导入之前，这类框架在第一次被导入的时候就会读取CUDA环境变量初始化GPU上下文，你导完框架再改这个变量完全没用。
设置完可以跑几行代码验证是否生效：

• PyTorch验证：

import torch
print(torch.cuda.device_count()) # 输出值应该等于你指定的GPU数量

• TensorFlow验证：

import tensorflow as tf
print(len(tf.config.list_physical_devices('GPU'))) # 输出值匹配你指定的卡数就对了

注意你指定的物理卡会被框架重新从0开始编号，比如你选了物理3号卡，框架里看到的cuda:0对应的就是物理3号卡，属于正常现象，不是设置错了。

3. 启动命令行内直接指定 #

嫌提前export麻烦的话，直接把环境变量写在启动命令前面就行，效果和第一种方法完全一致：

CUDA_VISIBLE_DEVICES=3 jupyter notebook

常见避坑点 #

• 已经跑过GPU代码的Notebook，临时改os.environ里的CUDA_VISIBLE_DEVICES没用，框架已经完成GPU初始化了，必须重启Notebook内核，再在最开头设置才生效。
• 每次选卡前都要先跑nvidia-smi确认占用，别靠记忆选卡，容易误占别人的任务资源。
• 要是你给Jupyter配了多套独立内核，方法1设置的全局变量对所有新开的内核都有效，如果单个Notebook要换卡，用方法2单独设置再重启对应内核就行。

内容的提问来源于stack exchange，提问作者Arist12