内存管理#

CuPy 默认使用内存池进行内存分配。内存池通过减轻内存分配和 CPU/GPU 同步的开销，显著提高了性能。

CuPy 中有两种不同的内存池

设备内存池（GPU 设备内存），用于 GPU 内存分配。
固定内存池（不可交换的 CPU 内存），用于 CPU 到 GPU 的数据传输。

注意

当您监控内存使用情况时（例如，使用 nvidia-smi 查看 GPU 内存或使用 ps 查看 CPU 内存），您可能会注意到即使数组实例超出作用域后，内存仍未释放。这是预期行为，因为默认的内存池会“缓存”已分配的内存块。

有关内存管理 API 的详细信息，请参阅低级 CUDA 支持。

为了更方便地使用固定内存，我们还在 cupyx 命名空间中提供了一些高级 API，包括 cupyx.empty_pinned()、cupyx.empty_like_pinned()、cupyx.zeros_pinned() 和 cupyx.zeros_like_pinned()。它们返回由固定内存支持的 NumPy 数组。如果使用了 CuPy 的固定内存池，则固定内存会从内存池中分配。

注意

CuPy v8 及更高版本提供了一个FFT 计划缓存，如果在使用了 FFT 及相关函数，它可能会占用部分设备内存。占用的内存可以通过缩小或禁用缓存来释放。

内存池操作#

内存池实例提供了内存分配的统计信息。要访问默认内存池实例，请使用 cupy.get_default_memory_pool() 和 cupy.get_default_pinned_memory_pool()。您还可以释放内存池中持有的所有未使用内存块。详细信息请参阅下面的示例代码

import cupy
import numpy

mempool = cupy.get_default_memory_pool()
pinned_mempool = cupy.get_default_pinned_memory_pool()

# Create an array on CPU.
# NumPy allocates 400 bytes in CPU (not managed by CuPy memory pool).
a_cpu = numpy.ndarray(100, dtype=numpy.float32)
print(a_cpu.nbytes)                      # 400

# You can access statistics of these memory pools.
print(mempool.used_bytes())              # 0
print(mempool.total_bytes())             # 0
print(pinned_mempool.n_free_blocks())    # 0

# Transfer the array from CPU to GPU.
# This allocates 400 bytes from the device memory pool, and another 400
# bytes from the pinned memory pool.  The allocated pinned memory will be
# released just after the transfer is complete.  Note that the actual
# allocation size may be rounded to larger value than the requested size
# for performance.
a = cupy.array(a_cpu)
print(a.nbytes)                          # 400
print(mempool.used_bytes())              # 512
print(mempool.total_bytes())             # 512
print(pinned_mempool.n_free_blocks())    # 1

# When the array goes out of scope, the allocated device memory is released
# and kept in the pool for future reuse.
a = None  # (or `del a`)
print(mempool.used_bytes())              # 0
print(mempool.total_bytes())             # 512
print(pinned_mempool.n_free_blocks())    # 1

# You can clear the memory pool by calling `free_all_blocks`.
mempool.free_all_blocks()
pinned_mempool.free_all_blocks()
print(mempool.used_bytes())              # 0
print(mempool.total_bytes())             # 0
print(pinned_mempool.n_free_blocks())    # 0

详细信息请参阅 cupy.cuda.MemoryPool 和 cupy.cuda.PinnedMemoryPool。

限制 GPU 内存使用#

您可以通过使用 CUPY_GPU_MEMORY_LIMIT 环境变量来硬性限制可分配的 GPU 内存量（详情请参阅环境变量）。

# Set the hard-limit to 1 GiB:
#   $ export CUPY_GPU_MEMORY_LIMIT="1073741824"

# You can also specify the limit in fraction of the total amount of memory
# on the GPU. If you have a GPU with 2 GiB memory, the following is
# equivalent to the above configuration.
#   $ export CUPY_GPU_MEMORY_LIMIT="50%"

import cupy
print(cupy.get_default_memory_pool().get_limit())  # 1073741824

您还可以使用 cupy.cuda.MemoryPool.set_limit() 来设置限制（或覆盖通过环境变量指定的值）。通过这种方式，您可以为每个 GPU 设备设置不同的限制。

import cupy

mempool = cupy.get_default_memory_pool()

with cupy.cuda.Device(0):
    mempool.set_limit(size=1024**3)  # 1 GiB

with cupy.cuda.Device(1):
    mempool.set_limit(size=2*1024**3)  # 2 GiB

注意

CUDA 会在内存池之外分配一些 GPU 内存（例如 CUDA 上下文、库句柄等）。根据使用情况，这些内存可能占用一到几百 MiB。这部分内存将不计入限制。

更改内存池#

您可以通过将内存分配函数传递给 cupy.cuda.set_allocator() / cupy.cuda.set_pinned_memory_allocator()，来使用您自己的内存分配器代替默认内存池。内存分配函数应接受 1 个参数（请求的大小，以字节为单位）并返回 cupy.cuda.MemoryPointer / cupy.cuda.PinnedMemoryPointer。

CuPy 提供了两个这样的分配器，用于在 GPU 上使用托管内存（managed memory）和流有序内存（stream ordered memory），详情请分别参阅 cupy.cuda.malloc_managed() 和 cupy.cuda.malloc_async()。要启用由托管内存支持的内存池，您可以构建一个新的 MemoryPool 实例，并将其分配器设置为 malloc_managed()，如下所示

import cupy

# Use managed memory
cupy.cuda.set_allocator(cupy.cuda.MemoryPool(cupy.cuda.malloc_managed).malloc)

请注意，如果您直接将 malloc_managed() 传递给 set_allocator() 而不构建 MemoryPool 实例，则内存释放时会立即返回给系统，这可能不是您期望的行为。

流有序内存分配器（Stream Ordered Memory Allocator）是 CUDA 11.2 之后添加的新功能。CuPy 提供了一个实验性的接口。与 CuPy 的内存池类似，流有序内存分配器也会以流有序的方式从内存池中异步地分配/释放内存。主要区别在于它是 NVIDIA 在 CUDA 驱动中实现的内置功能，因此同一进程中的其他 CUDA 应用程序可以轻松地从同一内存池分配内存。

要启用管理流有序内存的内存池，您可以构建一个新的 MemoryAsyncPool 实例

import cupy

# Use asynchronous stream ordered memory
cupy.cuda.set_allocator(cupy.cuda.MemoryAsyncPool().malloc)

# Create a custom stream
s = cupy.cuda.Stream()

# This would allocate memory asynchronously on stream s
with s:
    a = cupy.empty((100,), dtype=cupy.float64)

请注意，在这种情况下，我们不使用 MemoryPool 类。MemoryAsyncPool 接受与 MemoryPool 不同的输入参数，用于指示使用哪个内存池。更多详细信息请参阅 MemoryAsyncPool 的文档。

请注意，如果您直接将 malloc_async() 传递给 set_allocator() 而不构建 MemoryAsyncPool 实例，则将使用设备的当前内存池。

使用流有序内存时，重要的是您自己维护正确的流语义，例如使用 Stream 和 Event API（详情请参阅流和事件）；CuPy 不会试图为您智能地处理。在释放时，内存会异步释放，优先在分配该内存的流上释放（第一次尝试），或在任何当前的 CuPy 流上释放（第二次尝试）。允许在分配内存的流上的所有内存释放之前销毁该流。

此外，内部使用 cudaMalloc（CUDA 默认的同步分配器）的应用程序/库可能与流有序内存分配器产生意想不到的交互。具体来说，释放到内存池的内存可能无法立即被 cudaMalloc 看到，从而可能导致内存不足错误。在这种情况下，您可以调用 free_all_blocks()，或者手动执行（事件/流/设备）同步，然后重试。

目前，MemoryAsyncPool 接口是实验性的。特别是，虽然其 API 与 MemoryPool 基本相同，但由于 CUDA 的限制，该内存池的几个方法需要足够新的驱动（当然，还需要支持的硬件、CUDA 版本和平台）。

您甚至可以使用以下代码禁用默认内存池。请确保在进行任何其他 CuPy 操作之前执行此操作。

import cupy

# Disable memory pool for device memory (GPU)
cupy.cuda.set_allocator(None)

# Disable memory pool for pinned memory (CPU).
cupy.cuda.set_pinned_memory_allocator(None)