PyTorch 中通过类 torch.utils.data.DataLoader 提供给模型训练数据集。该类中有一个参数 pin_memory 表示是否采用锁页内存。那什么是锁页内存，它有什么用？本篇对其进行介绍，以帮助大家了解其中的原理，设置合适的参数值，加快模型的训练。

什么是锁页内存

我们都知道计算机除了 CPU 进行计算外，还离不开内存，因为内存中数据读取、交换比较快，所以，计算机程序运行时，将中间数据存放到内存中将大大提高程序的运行速度。但内存相对比较昂贵，因此，计算机的内存往往容量有限，在 Linux 中，常见有虚拟内存（通过 htop 命令可以查看 Swp 就是虚拟内存），其是在硬盘的每个区域划分出来的模拟内存的一块区域，在内存空间不足时，可以作为临时的内存空间使用，但虚拟内存毕竟是在硬盘上，数据的读取速度明显比内存慢。

在内存中，数据的存储分为两类，一类是锁页，一类不是锁页。锁页内存存放的数据在任何情况下都不会与虚拟内存进行数据交换。而不锁页内存当内存不足时，数据会存放在虚拟内置中。值得一提的是，GPU 的显存全部都是锁页内存。

设置 pin_memory

了解了什么是锁页内存，在编写代码时，就可以有针对性的设置 torch.utils.data.DataLoader 中的参数 pin_memory 。

当计算机的内存充足时，可设置 pin_memory=True；如果为 True，数据加载器将在返回之前将张量复制到 CUDA 固定内存中。
当计算集的内存不足时，需设置 pin_memory=False.

在 PyTorch 中，因为 pin_memory 与电脑硬件有关，开发者不能确保每一个深度学习研究员都有高端的计算设备，因此，pin_memory 的默认值为 False.

如果不是很清楚是否应该设置 pin_memory 为 True，可以通过如下的代码测试：

import multiprocessing
import time

use_cuda = torch.cuda.is_available()
core_number = multiprocessing.cpu_count()
batch_size = 64
best_num_worker = [0, 0]
best_time = [99999999, 99999999]
print("cpu_count =", core_number)
    
def loading_time(num_workers, pin_memory):
        kwargs = {"num_workers": num_workers, "pin_memory": pin_memory} if use_cuda else {}
        train_loader = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=args.batch_size,
            shuffle=(train_sampler is None),
            sampler=train_sampler,
            **kwargs
        )
        start = time.time()
        for epoch in range(4):
            for batch_idx, (data, target) in enumerate(train_loader):
                if batch_idx == 15:
                    break
        end = time.time()
        print("Used {} second with num_workers = {}".format(end - start, num_workers))
        return end - start


    for pin_memory in [False, True]:
        print("While pin_memory =", pin_memory)
        for num_workers in range(0, core_number * 2 + 1, 4):
            current_time = loading_time(num_workers, pin_memory)
            if current_time < best_time[pin_memory]:
                best_time[pin_memory] = current_time
                best_num_worker[pin_memory] = num_workers
            else:  # assuming its a convex function
                if best_num_worker[pin_memory] == 0:
                    the_range = []
                else:
                    the_range = list(
                        range(best_num_worker[pin_memory] - 3, best_num_worker[pin_memory])
                    )
                for num_workers in the_range + list(
                    range(best_num_worker[pin_memory] + 1, best_num_worker[pin_memory] + 4)
                ):
                    current_time = loading_time(num_workers, pin_memory)
                    if current_time < best_time[pin_memory]:
                        best_time[pin_memory] = current_time
                        best_num_worker[pin_memory] = num_workers
                        break
    if best_time[0] < best_time[1]:
        print("Best num_workers =", best_num_worker[0], "with pin_memory = False")
    else:
        print("Best num_workers =", best_num_worker[1], "with pin_memory = True")
    return

将这些代码放到：GitHub-pytorch-examples-imagenet-main.py 中的 main_worker 函数中的下面位置：

    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
        num_workers=args.workers, pin_memory=True, sampler=train_sampler)
 
---------------- here ---------------- 
 
    val_loader = torch.utils.data.DataLoader(
        datasets.ImageFolder(valdir, transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            normalize,
        ])),
        batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=True)