锁页内存与异步拷贝

Pinned Memory

理论上其通过将数据放在锁页内存，且通过异步传输可以加速数据从内存到显存的传递速率

其案例：

import os
import time
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader


# ===== 自动生成 CSV 文件 =====
csv_path = r"E:\data\data.csv"
if not os.path.exists(csv_path):
    N = 1000000  # 行数，可以调大
    M = 21     # 列数，最后一列作为 label
    df = pd.DataFrame(np.random.randn(N, M))
    df.to_csv(csv_path, index=False)
    print(f"已生成测试文件 {csv_path}")


# ===== 自定义 Dataset =====
class CSVDataset(Dataset):
    def __init__(self, csv_path):
        self.data = pd.read_csv(csv_path)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        row = self.data.iloc[idx]
        x = torch.tensor(row.iloc[:-1].values, dtype=torch.float32)
        y = torch.tensor(row.iloc[-1], dtype=torch.float32)  # ✅ 改这里
        return x, y


# ===== 测试函数 =====
def run_dataloader(csv_path, pin_memory=False):
    dataset = CSVDataset(csv_path)

    dataloader = DataLoader(
        dataset,
        batch_size=1024,
        shuffle=True,
        num_workers=2,
        pin_memory=pin_memory # 锁页内存,固定放在某个页面,可加速数据拷贝
    )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    start = time.time()
    for data, target in dataloader:
        # 关键：用 non_blocking=True
        data = data.to(device, non_blocking=True) # 异步传输
        target = target.to(device, non_blocking=True) # 异步传输
        # 模拟训练计算
        _ = data * 2.0 + 1.0
    end = time.time()

    print(f"pin_memory={pin_memory}, time={end-start:.4f} s")


if __name__ == "__main__":
    run_dataloader(csv_path, pin_memory=False)
    run_dataloader(csv_path, pin_memory=True)

不过就结果而言，基本没带来提升，甚至很多情况下还带来负面反馈，建议谨慎使用

未来有待进一步研究的点：

1.定位实验中未能展现效果的原因并总结影响因素

2.利用该方案成功实现正向效果