如何直接在共享内存中创建大尺寸PyTorch CPU张量？

需要创建多GB级的PyTorch CPU张量，内存仅能容纳该张量的一份副本。由于后续要给设置了num_workers > 0的DataLoader多进程工作线程使用，必须将其设为共享内存张量，且不能出现内存中同时存在两份副本的情况。

已尝试的方法及问题 #

1. 直接创建张量后使用DataLoader #

代码：

v = torch.empty(25 * 2**30)
loader = DataLoader(dataset, num_workers=2, persistent_workers=True)  
# dataset在__init__中使用v

问题：触发Couldn't open shared file mapping:..错误，原因是张量先在私有内存创建，再复制到共享内存，两份副本超出内存容量。

2. 显式调用share_memory_() #

代码：

v = torch.empty(25 * 2**30)
v.share_memory_()

问题：和方法1同理，先创建私有内存张量再复制到共享内存，导致内存中存在两份副本，触发相同错误。

3. 使用torch.from_file #

可行但不符合需求：需要将大张量写入磁盘，速度慢，且依赖磁盘IO，不符合内存直接分配的要求。

4. TorchStore模块 #

尚未正式纳入PyTorch，无法在稳定版本中使用。

5. 使用multiprocessing.shared_memory.SharedMemory #

代码：

import torch
import numpy as np
from multiprocessing import shared_memory
import psutil

print(f"Available memory: {psutil.virtual_memory().available / (1024 ** 3):.2f} GB")
print(f"Total memory: {psutil.virtual_memory().total / (1024 ** 3):.2f} GB")

tensor_shape = (2**30, 25)
dtype = np.float32
num_elements = np.prod(tensor_shape)
size = int(num_elements * np.dtype(dtype).itemsize)
print(f"Requesting allocation of {size / 2 ** 30:.2f} GB")

sh_mem = shared_memory.SharedMemory(create=True, size=size)
np_array = np.ndarray(tensor_shape, dtype=dtype, buffer=sh_mem.buf)
tensor = torch.from_numpy(np_array)
tensor.fill_(0)
print(f"Allocated {size / 2**30:.2f} GB")

问题：当申请内存超过总内存一半时报错（Windows 10和Ubuntu 24.04表现一致），错误信息：

Available memory: 176.63 GB
Total memory: 191.87 GB
Requesting allocation of 100.00 GB
Traceback (most recent call last):
  File "D:\Sci\NetOnset\CheckShared.py", line 22, in 
    sh_mem = shared_memory.SharedMemory(create=True, size=size)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Miniconda3\envs\cuda124\Lib\multiprocessing\shared_memory.py", line 151, in __init__
    self._mmap = mmap.mmap(-1, size, tagname=temp_name)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
OSError: [WinError 1455] The paging file is too small for this operation to complete

核心问题 #

是否仍存在无需复制即可在共享内存中创建大尺寸PyTorch张量的方法？

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