PyCUDA与TensorFlow能否协同使用？启动TF后PyCUDA代码崩溃求助

当然可以让PyCUDA和TensorFlow协同使用啦！不过你遇到的这个cuFuncSetBlockShape failed: invalid resource handle错误，本质是两者在CUDA上下文管理上的冲突——TensorFlow初始化时会接管CUDA上下文，哪怕调用sess.close()，它也不会彻底释放或重置上下文，导致后续PyCUDA无法获取有效的资源句柄。

下面给你几个实用的解决思路：

1. 手动管理CUDA上下文（推荐） #

别用pycuda.autoinit（它会自动创建默认上下文，很容易和TensorFlow的上下文打架），而是手动获取TensorFlow的CUDA上下文让PyCUDA复用，或者手动创建并切换上下文。

举个代码例子：

import tensorflow as tf
import pycuda.driver as drv
import numpy as np
from pycuda.compiler import SourceModule

# 先初始化TensorFlow会话
sess = tf.Session()

# 获取TensorFlow正在使用的CUDA上下文
ctx = drv.Context.get_current()

# 现在可以正常跑PyCUDA代码了
mod = SourceModule("""
    __global__ void multiply_them(float *dest, float *a, float *b)
    {
        const int i = threadIdx.x;
        dest[i] = a[i] * b[i];
    }
""")

multiply_them = mod.get_function("multiply_them")

a = np.random.randn(400).astype(np.float32)
b = np.random.randn(400).astype(np.float32)
dest = np.zeros_like(a)

multiply_them(
    drv.Out(dest), drv.In(a), drv.In(b),
    block=(400,1,1), grid=(1,1)
)

# 用完后记得清理
sess.close()
ctx.pop()

2. 让PyCUDA先创建上下文 #

如果你的业务逻辑是先跑PyCUDA再用TensorFlow，那可以先手动初始化PyCUDA的上下文，让TensorFlow后续复用这个上下文：

import pycuda.driver as drv
import tensorflow as tf
import numpy as np
from pycuda.compiler import SourceModule

# 手动初始化PyCUDA上下文
drv.init()
dev = drv.Device(0)  # 选择第0块GPU
ctx = dev.make_context()

# 先执行PyCUDA代码
mod = SourceModule("""
    __global__ void multiply_them(float *dest, float *a, float *b)
    {
        const int i = threadIdx.x;
        dest[i] = a[i] * b[i];
    }
""")

multiply_them = mod.get_function("multiply_them")

a = np.random.randn(400).astype(np.float32)
b = np.random.randn(400).astype(np.float32)
dest = np.zeros_like(a)

multiply_them(
    drv.Out(dest), drv.In(a), drv.In(b),
    block=(400,1,1), grid=(1,1)
)

# 再初始化TensorFlow，它会自动复用当前的CUDA上下文
sess = tf.Session()
# 这里可以写你的TensorFlow逻辑...
sess.close()

# 最后清理PyCUDA上下文
ctx.pop()

3. 多GPU场景：隔离设备 #

如果你有多个GPU，可以给TensorFlow和PyCUDA分配不同的GPU，从根源上避免上下文冲突。

你可以在运行脚本前设置环境变量：

# 让TensorFlow用GPU 0，PyCUDA用GPU 1
CUDA_VISIBLE_DEVICES=0 python your_tensorflow_part.py
CUDA_VISIBLE_DEVICES=1 python your_pycuda_part.py

也可以在代码里直接指定：

import os
# 先给TensorFlow分配GPU 0
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import tensorflow as tf

# 再给PyCUDA分配GPU 1
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
import pycuda.driver as drv

核心思路就是：明确CUDA上下文的归属，避免自动初始化带来的冲突，要么复用同一个上下文，要么隔离不同设备。

内容的提问来源于stack exchange，提问作者Chi-Fang Hsieh