我一直在尝试用一些cuda库为python做线性代数运算。到目前为止,最直接的似乎是Cupy。
然而,我在使用这些库时遇到的问题是,每当运算量大到足以让gpu计算感兴趣时,我就会遇到内存错误。
在Cupy或类似的库中是否有一些方法可以自动进行内存管理?
你可以尝试要求CuPy在进入内存密集区域之前回收所有未使用的GPU内存。
mempool = cupy.get_default_memory_pool()
mempool.free_all_blocks()
任何CuPy数组在调用前的ref数为0的,都会被deallocated,底层GPU内存会被返回到内存池中。