高性能PyTorch是如何炼成的？整理的10条脱坑指南

如何用最少的精力，完成最高效的 PyTorch 训练?一位有着 PyTorch 两年使用经历的 Medium 博主最近分享了他在这方面的 10 个真诚建议。

在 Efficient PyTorch 这一部分中，作者提供了一些识别和消除 I/O 和 CPU 瓶颈的技巧。第二部分阐述了一些高效张量运算的技巧，第三部分是在高效模型上的 debug 技巧。

在阅读这篇文章之前，你需要对 PyTorch 有一定程度的了解。

好吧，从最明显的一个开始：

建议 0：了解你代码中的瓶颈在哪里

命令行工具比如 nvidia-smi、htop、iotop、nvtop、py-spy、strace 等，应该成为你最好的伙伴。你的训练管道是否受 CPU 约束?IO 约束?GPU 约束?这些工具将帮你找到答案。

这些工具你可能从未听过，即使听过也可能没用过。没关系。如果你不立即使用它们也可以。只需记住，其他人可能正在用它们来训练模型，速度可能会比你快 5%、10%、15%-…… 最终可能会导致面向市场或者工作机会时候的不同结果。

数据预处理

几乎每个训练管道都以 Dataset 类开始。它负责提供数据样本。任何必要的数据转换和扩充都可能在此进行。简而言之，Dataset 能报告其规模大小以及在给定索引时，给出数据样本。

如果你要处理类图像的数据(2D、3D 扫描)，那么磁盘 I/O 可能会成为瓶颈。为了获取原始像素数据，你的代码需要从磁盘中读取数据并解码图像到内存。每个任务都是迅速的，但是当你需要尽快处理成百上千或者成千上万个任务时，可能就成了一个挑战。像 NVidia 这样的库会提供一个 GPU 加速的 JPEG 解码。如果你在数据处理管道中遇到了 IO 瓶颈，这种方法绝对值得一试。

还有另外一个选择，SSD 磁盘的访问时间约为 0.08–0.16 毫秒。RAM 的访问时间是纳秒级别的。我们可以直接将数据存入内存。

建议 1：如果可能的话，将数据的全部或部分移至 RAM。

如果你的内存中有足够多的 RAM 来加载和保存你的训练数据，这是从管道中排除最慢的数据检索步骤最简单的方法。

这个建议可能对云实例特别有用，比如亚马逊的 p3.8xlarge。该实例有 EBS 磁盘，它的性能在默认设置下非常受限。但是，该实例配备了惊人的 248Gb 的 RAM。这足够将整个 ImageNet 数据集存入内存了!你可以通过以下方法达到这一目标：

class RAMDataset(Dataset):   def __init__(image_fnames, targets):     self.targets = targets     self.images = []     for fname in tqdm(image_fnames, desc="Loading files in RAM"):       with open(fname, "rb") as f:         self.images.append(f.read())    def __len__(self):     return len(self.targets)    def __getitem__(self, index):     target = self.targets[index]     image, retval = cv2.imdecode(self.images[index], cv2.IMREAD_COLOR)     return image, target 

我个人也面对过这个瓶颈问题。我有一台配有 4x1080Ti GPUs 的家用 PC。有一次，我采用了有 4 个 NVidia Tesla V100 的 p3.8xlarge 实例，然后将我的训练代码移到那里。鉴于 V100 比我的 oldie 1080Ti 更新更快的事实，我期待看到训练快 15–30%。出乎意料的是，每个时期的训练时间都增加了。这让我明白要注意基础设施和环境差异，而不仅仅是 CPU 和 GPU 的速度。

根据你的方案，你可以将每个文件的二进制内容保持不变，并在 RAM 中进行即时解码，或者对未压缩的图像进行讲解码，并保留原始像素。但是无论你采用什么方法，这里有第二条建议：

建议 2：解析、度量、比较。每次你在管道中提出任何改变，要深入地评估它全面的影响。

假设你对模型、超参数和数据集等没做任何改动，这条建议只关注训练速度。你可以设置一个魔术命令行参数(魔术开关)，在指定该参数时，训练会在一些合理的数据样例上运行。利用这个特点，你可以迅速解析管道。

# Profile CPU bottlenecks python -m cProfile training_script.py --profiling  # Profile GPU bottlenecks nvprof --print-gpu-trace python train_mnist.py  # Profile system calls bottlenecks strace -fcT python training_script.py -e trace=open,close,read  Advice 3: *Preprocess everything offline* 

建议 3：离线预处理所有内容

如果你要训练由多张 2048x2048 图像制成的 512x512 尺寸图像，请事先调整。如果你使用灰度图像作为模型的输入，请离线调整颜色。如果你正在进行自然语言处理(NLP)，请事先做分词处理(tokenization)，并存入磁盘。在训练期间一次次重复相同的操作没有意义。在进行渐进式学习时，你可以以多种分辨率保存训练数据的，这还是比线上调至目标分辨率更快。

对于表格数据，请考虑在创建 Dataset 时将 pd.DataFrame 目标转换为 PyTorch 张量。

建议 4：调整 DataLoader 的工作程序

PyTorch 使用一个 DataLoader 类来简化用于训练模型的批处理过程。为了加快速度，它可以使用 Python 中的多进程并行执行。大多数情况下，它可以直接使用。还有几点需要记住：

每个进程生成一批数据，这些批通过互斥锁同步可用于主进程。如果你有 N 个工作程序，那么你的脚本将需要 N 倍的 RAM 才能在系统内存中存储这些批次的数据。具体需要多少 RAM 呢?

我们来计算一下：

（编辑：晋中站长网）

【声明】本站内容均来自网络，其相关言论仅代表作者个人观点，不代表本站立场。若无意侵犯到您的权利，请及时与联系站长删除相关内容!