大尺寸视频转torch.tensor内存不足解决方案咨询

1. 内存映射张量（Memory-Mapped Tensor） #

这是适配你场景的最优方案——无需一次性将全量数据加载进RAM，通过磁盘内存映射实现按需读写，最终生成单一.pt文件。

步骤实现： #

（1）预计算尺寸并创建内存映射张量

优先使用uint8数据类型（对应视频原始像素的0-255范围），能最大化压缩磁盘占用：

import torch

# 张量形状定义：(帧数, 通道数, 高度, 宽度)，假设视频为RGB三通道
tensor_shape = (20000, 3, 480, 854)
dtype = torch.uint8

# 创建内存映射文件，mode='w+'支持读写操作
mmap_tensor = torch.mmap_file(
    'video_full_tensor.pt',
    shape=tensor_shape,
    dtype=dtype,
    mode='w+'
)

（2）逐帧读取视频并写入映射张量

用OpenCV或PyAV逐帧读取视频，直接写入内存映射张量的对应索引（操作仅占用单帧内存）：

import cv2

cap = cv2.VideoCapture('your_input_video.mp4')
frame_idx = 0

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    # 转换OpenCV默认的BGR格式为RGB，调整维度为(3, H, W)
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    frame_tensor = torch.from_numpy(frame_rgb).permute(2, 0, 1).to(dtype=dtype)
    # 写入内存映射张量（直接操作磁盘，不占满RAM）
    mmap_tensor[frame_idx] = frame_tensor
    frame_idx += 1

cap.release()

（3）加载与使用

后续加载该.pt文件时，数据不会一次性涌入RAM，仅在访问特定帧时触发磁盘读取：

# 加载内存映射张量
loaded_tensor = torch.load('video_full_tensor.pt', map_location='cpu')
# 访问第100帧时才会读取对应磁盘数据
target_frame = loaded_tensor[100]

2. 压缩版内存映射方案 #

若想进一步降低磁盘占用，可在内存映射基础上配合PyTorch的压缩序列化机制：

# 写入完成后，用压缩格式保存张量对象
torch.save(mmap_tensor, 'video_full_tensor_compressed.pt', _use_new_zipfile_serialization=True)

加载时仍支持按需读取，磁盘占用比未压缩版本低30%-50%（取决于视频内容）。

关键优化细节 #

• 强制使用uint8类型：若改用float32存储，磁盘占用会直接放大4倍，完全没必要（原始视频像素本身是8位无符号整数）。
• 高效视频读取工具：长视频场景下，PyAV比OpenCV读取效率更高，替换示例：

import av

container = av.open('your_input_video.mp4')
frame_idx = 0

for frame in container.decode(video=0):
    frame_rgb = frame.to_rgb().to_ndarray()
    frame_tensor = torch.from_numpy(frame_rgb).permute(2, 0, 1).to(dtype=dtype)
    mmap_tensor[frame_idx] = frame_tensor
    frame_idx += 1

• 跨平台注意事项：内存映射文件在Windows与Linux/macOS间存在格式差异，建议在同一系统环境下完成写入和读取。

内容的提问来源于stack exchange，提问作者K.XGai