形状对GPU上的张量流的内存分配的影响

Question

当我在具有tensorflow的nvidia GPU上创建形状为(L, N, N)的变量时，该变量是否占据L个正方形块的平铺，其中每个块为N x N？

例如，将三个正方形的内存块平铺在一起将是：

# --- #
| NxN |
# --- #
| NxN |
# --- #
| NxN |
# --- #

张量流如何为形状为(a, b, N, N)的变量分配GPU内存？

Answer 1

GPU上的内存分配为一维字节数组，并尽可能对齐。在您的示例中，形状为(L, N, N)的张量将是大小为8*L*N*N个字节的一维数组。假设值为(a, b, N, N)，则8*a*b*N*N的大小为float64个字节。

您可以通过阅读TensorFlow源代码来找到答案：

• gpu_cudamalloc_allocator.h用GPUcudaMallocAllocator声明void* AllocateRaw(size_t alignment, size_t num_bytes)类，该类在
中实现
• gpu_cudamalloc_allocator.cc。
• allocator.h声明Allocator的类GPUcudaMallocAllocator，该类声明void* AllocateRaw(size_t alignment, size_t num_bytes)。

由于这些AllocateRaw函数采用单一大小，因此所有TensorFlow内存都在底层的1-D数组中分配。这在实践中很常见，因为它很容易将高阶数组表示形式映射到连续的一维内存块上，并且可以最大程度地提高向设备传输和向后传输的吞吐量。