ubuntu pytorch怎么共享内存

在 Ubuntu 上，**PyTorch 的“共享内存”**通常有两种常见含义，我分别说明，并给出用法和注意事项。

一、进程间共享内存（最常用：multiprocessing + Tensor 共享内存）

PyTorch 在 Linux/Ubuntu 上，利用 /dev/shm（tmpfs） 实现进程间张量共享。

1️⃣ 多进程（DataLoader）默认就是共享内存

from torch.utils.data import DataLoader

loader = DataLoader(
    dataset,
    batch_size=32,
    num_workers=4,   # 多进程
    pin_memory=False
)

• num_workers > 0 时，数据通过 共享内存 在 worker 和主进程间传递
• 不需要你手动管理

2️⃣ 手动共享 Tensor（多进程）

import torch
import torch.multiprocessing as mp

def worker(tensor):
    tensor.add_(1)

if __name__ == "__main__":
    tensor = torch.zeros(5)
    tensor.share_memory_()   # 关键

    p = mp.Process(target=worker, args=(tensor,))
    p.start()
    p.join()

    print(tensor)

✅ share_memory_()：

• 把 Tensor 放到共享内存
• 多进程可直接读写
• 仅支持 CPU Tensor

3️⃣ 共享内存大小不足（常见问题）

Ubuntu 默认 /dev/shm 可能太小（如 Docker 中）：

df -h /dev/shm

解决方法

sudo mount -o remount,size=4G /dev/shm

Docker：

docker run --shm-size=4g ...

二、GPU 显存共享（CUDA 共享）

1️⃣ 多进程共享同一块 GPU

• CUDA 本身支持
• PyTorch 默认支持多进程使用同一 GPU

tensor = torch.zeros(3, device="cuda")

⚠️ 注意：

• 不要同时写同一块显存
• 多进程训练一般用 DistributedDataParallel

2️⃣ 多进程 DDP（推荐）

torchrun --nproc_per_node=2 train.py

三、Python 层面共享内存（非 PyTorch 专用）

from multiprocessing import shared_memory

但 ❌ 不推荐用于 Tensor（序列化成本高）

四、常见坑总结

五、你可以告诉我

你是：

• ✅ DataLoader 太慢？
• ✅ 多进程推理？
• ✅ 多卡训练？
• ✅ 想省显存？

我可以给你具体代码模板。