cuML tSNE批量处理异常：无报错但无法正常执行

环境信息 #

• 系统：Windows笔记本 + WSL2
• 硬件：GTX 4050笔记本GPU、AMD Ryzen 7 7840HS
• 环境创建命令：

conda create --solver=libmamba -n rapids-24.02 -c rapidsai -c conda-forge -c nvidia  \
    rapids=24.02 python=3.10 cuda-version=12.0 \
    jupyterlab tensorflow pytorch

• 单独运行tSNE测试代码正常，整合到批量处理流程后无报错但无法正常执行，PCA等其他cuML方法可正常运行。

排查方向与解决方法 #

1. 显存占用冲突 #

tSNE（尤其是barnes_hut方法）对显存需求较高，批量流程中其他步骤可能已占用大量显存，导致tSNE无法分配足够资源陷入等待。

• 排查：在tSNE执行前后添加显存监控：

import pynvml
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
print(f"剩余显存: {mem_info.free / 1024**2:.2f} MB")

• 解决：
  • tSNE执行前清理无用GPU张量，调用torch.cuda.empty_cache()或cudf.DataFrame.clear()释放显存；
  • 降低perplexity参数（默认30），或减少单批次样本量；
  • 切换为method='fft'（RAPIDs 23.08+支持），该方法显存占用更低。

2. 数据类型/格式不匹配 #

单独测试用numpy数组，批量流程中若传入pandas DataFrame而非cudf对象，会触发cuML tSNE的CPU fallback，运行极慢看似无响应。

• 排查：打印输入数据的类型与设备：

print(f"输入数据类型: {type(df)}")
if hasattr(df, 'device'):
    print(f"数据所在设备: {df.device}")

• 解决：确保传入tSNE的是cudf DataFrame或cuPy数组：

# 转换为cudf DataFrame
df_cudf = cudf.DataFrame(df)
# 或转换为cuPy数组
import cupy as cp
df_cupy = cp.array(df)

3. 线程/异步执行冲突 #

批量流程若使用多线程/异步框架，可能与RAPIDs的CUDA上下文管理冲突，导致tSNE执行被阻塞。

• 排查：暂时移除多线程/异步逻辑，以单线程运行tSNE模块验证。
• 解决：
  • 在主线程初始化CUDA上下文，避免子线程调用cuML方法；
  • 多进程场景下，每个进程单独初始化RAPIDs环境，禁止上下文共享。

4. WSL2 CUDA资源限制 #

WSL2默认的GPU调度限制可能导致批量处理时资源竞争，tSNE无法获取足够算力。

• 解决：
  • 升级WSL2到最新版本，确保CUDA驱动与RAPIDs指定的cuda-version=12.0匹配；
  • 在WSL2中设置环境变量绑定GPU：

export CUDA_VISIBLE_DEVICES=0

• 更新Windows主机的GPU驱动，避免跨环境兼容性问题。

5. tSNE迭代参数设置 #

批量流程中若n_iter过高或learning_rate设置不当，会导致tSNE运行时间过长，看似无响应。

• 解决：
  • 初始化tSNE时开启verbose=True查看迭代进度：

tsne = TSNE(n_components=2, method='barnes_hut', random_state=42, verbose=True)

• 临时降低n_iter（如设为200）验证能否快速完成，再逐步调整到合理值。

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