Python多线程IO读数据后GPU处理：如何避免多线程抢占GPU？

问题背景 #

我用多线程读取文件获取数据，之后调用GPU加速函数处理，但系统只有一块GPU，必须保证同一时间只有一个线程访问GPU。

当前采用单线程池实现，代码如下：

from multiprocessing.dummy import Pool as ThreadPool

def my_fun():
    urls = [url_1, url_2,....,url_n]
    pool = ThreadPool(1) # <--- 单线程！
    pool.map(read_and_process, urls)
    pool.close()
    pool.join()

def read_and_process(url):
    data = read(url) # IO密集型操作
    process(data)    # GPU处理操作

但单线程池和普通循环效率一致，完全发挥不了多线程的优势；要是增加线程数，程序就会崩溃。

我的需求：

• 利用多线程处理IO密集的读数据操作
• 避免多线程抢占GPU资源
• 由于单条数据超过600MB，总共有数千条，无法把所有数据存入内存后统一处理，所以process()不能移出read_and_process函数。

我尝试用线程锁修改了代码，目前看起来能运行，想请大家验证并给出反馈，修改后的代码如下：

from multiprocessing.dummy import Pool as ThreadPool
from multiprocessing.dummy import Lock

def my_fun():
    urls = [url_1, url_2,....,url_n]
    lock = Lock()
    items = [(d, lock) for d in urls]

    with ThreadPool(4):
        pool.starmap_async(read_and_process, items)
    

def read_and_process(url, lock):
    data = read(url) # IO密集型操作
    with lock:
        process(data)    # GPU处理操作

验证与反馈 #

你的思路完全正确，用线程锁控制GPU操作的串行执行，同时让IO读取部分并行，刚好契合IO密集+GPU计算的场景，以下是具体验证点和优化建议：

正确性验证 #

• 线程锁Lock能保证同一时间只有一个线程进入with lock代码块，也就是process(data)会被串行执行，不会出现多线程抢占GPU的情况，从逻辑上彻底解决了多线程访问GPU崩溃的问题。
• multiprocessing.dummy的Lock是线程锁，和线程池配套使用完全正确——因为dummy.Pool本质是基于线程的池，不是进程池，线程锁能正常生效。

代码中的小问题 #

1. 线程池上下文管理器使用错误：
   你用with ThreadPool(4):时没有把池对象赋值给pool，后续调用pool.starmap_async会直接报错。正确写法：

   with ThreadPool(4) as pool:
       pool.starmap_async(read_and_process, items)
   

   另外，starmap_async是异步调用，如果你需要等待所有任务完成后再执行后续逻辑，要加上pool.close()和pool.join()，或者改用同步版本的starmap，否则主线程可能提前退出，导致子线程被强制终止。

2. 锁的传递可以优化：
   不需要把锁和每个URL打包成元组传递，可通过functools.partial把锁绑定到函数参数上，代码更简洁：

   from functools import partial
   
   def my_fun():
       urls = [url_1, url_2,....,url_n]
       lock = Lock()
       # 绑定锁到函数参数
       bound_func = partial(read_and_process, lock=lock)
       
       with ThreadPool(4) as pool:
           pool.map(bound_func, urls)
           pool.close()
           pool.join()
   
   def read_and_process(url, lock):
       data = read(url)
       with lock:
           process(data)
   

额外优化建议 #

• 线程池大小调整：IO密集型任务的线程池大小可以适当调大（比如8-16，根据CPU核心数和网络情况），因为IO操作大部分时间处于等待状态，多线程能更充分利用CPU空闲时间。
• GPU内存释放：process(data)执行完成后，建议手动释放数据占用的GPU内存（比如PyTorch的del data+torch.cuda.empty_cache()，TensorFlow的对应清理操作），避免GPU内存累积溢出，尤其是处理大尺寸数据时。
• 异常处理：在read和process中加入异常捕获，避免单个任务失败导致整个线程池挂掉，示例：

  def read_and_process(url, lock):
      try:
          data = read(url)
          with lock:
              process(data)
      except Exception as e:
          print(f"处理URL {url}失败: {str(e)}")
  

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