使用自定义 PyTorch 运算符优化深度学习数据输入管道

简介: 使用自定义 PyTorch 运算符优化深度学习数据输入管道

这篇文章中,我们讨论 PyTorch 对创建自定义运算符的支持,并演示它如何帮助我们解决数据输入管道的性能瓶颈、加速深度学习工作负载并降低训练成本。

构建 PyTorch 扩展

PyTorch 提供了多种创建自定义操作的方法,包括使用自定义模块和/或函数扩展 torch.nn。在这篇文章中,我们感兴趣的是 PyTorch 对集成定制 C++ 代码的支持。此功能很重要,因为某些操作在 C++ 中比在 Python 中更有效和/或更容易地实现。使用指定的 PyTorch 实用程序(例如 CppExtension),可以轻松地将这些操作作为 PyTorch 的“扩展”包含在内,而无需拉取和重新编译整个 PyTorch 代码库。由于我们对这篇文章的兴趣是加速基于 CPU 的数据预处理管道,因此我们只需使用 C++ 扩展即可,不需要 CUDA 代码。

玩具示例

在我们之前的文章中,我们定义了一个数据输入管道,首先解码 533x800 JPEG 图像,然后提取随机的 256x256 裁剪,经过一些额外的转换后,将其输入训练循环。我们使用 PyTorch Profiler 和 TensorBoard 来测量与从文件加载图像相关的时间,并承认解码的浪费。为了完整起见,我们复制以下代码:

import numpy as np
from PIL import Image
from torchvision.datasets.vision import VisionDataset
input_img_size = [533, 800]
img_size = 256

class FakeDataset(VisionDataset):
    def __init__(self, transform):
        super().__init__(root=None, transform=transform)
        size = 10000
        self.img_files = [f'{i}.jpg' for i in range(size)]
        self.targets = np.random.randint(low=0,high=num_classes,
                                         size=(size),dtype=np.uint8).tolist()

    def __getitem__(self, index):
        img_file, target = self.img_files[index], self.targets[index]
        img = Image.open(img_file)
        if self.transform is not None:
            img = self.transform(img)
        return img, target

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


transform = T.Compose(
    [T.PILToTensor(),
     T.RandomCrop(img_size),
     RandomMask(),
     ConvertColor(),
     Scale()])

据推测,如果我们能够仅解码我们感兴趣的作物,我们的管道会运行得更快。不幸的是,截至撰写本文时,PyTorch 不包含支持此功能的函数。然而,使用自定义操作创建工具,我们可以定义并实现我们自己的函数!

自定义 JPEG 图像解码和裁剪函数

libjpeg-turbo 库是一个 JPEG 图像编解码器,与 libjpeg 相比,它包含许多增强和优化。特别是,libjpeg-turbo 包含许多函数,使我们能够仅解码图像中的预定义裁剪,例如 jpeg_skip_scanlines 和 jpeg_crop_scanline。如果您在 conda 环境中运行,可以使用以下命令进行安装:

conda install -c conda-forge libjpeg-turbo

请注意,libjpeg-turbo 已预安装在我们将在下面的实验中使用的官方 AWS PyTorch 2.0 深度学习 Docker 映像中。
在下面的代码块中,我们修改了torchvision 0.15的decode_jpeg函数,以从输入的JPEG编码图像中解码并返回所请求的裁剪。

torch::Tensor decode_and_crop_jpeg(const torch::Tensor& data,
                                   unsigned int crop_y,
                                   unsigned int crop_x,
                                   unsigned int crop_height,
                                   unsigned int crop_width) {
  struct jpeg_decompress_struct cinfo;
  struct torch_jpeg_error_mgr jerr;

  auto datap = data.data_ptr<uint8_t>();
  // Setup decompression structure
  cinfo.err = jpeg_std_error(&jerr.pub);
  jerr.pub.error_exit = torch_jpeg_error_exit;
  /* Establish the setjmp return context for my_error_exit to use. */
  setjmp(jerr.setjmp_buffer);
  jpeg_create_decompress(&cinfo);
  torch_jpeg_set_source_mgr(&cinfo, datap, data.numel());

  // read info from header.
  jpeg_read_header(&cinfo, TRUE);

  int channels = cinfo.num_components;

  jpeg_start_decompress(&cinfo);

  int stride = crop_width * channels;
  auto tensor =
     torch::empty({int64_t(crop_height), int64_t(crop_width), channels},
                  torch::kU8);
  auto ptr = tensor.data_ptr<uint8_t>();

  unsigned int update_width = crop_width;
  jpeg_crop_scanline(&cinfo, &crop_x, &update_width);
  jpeg_skip_scanlines(&cinfo, crop_y);

  const int offset = (cinfo.output_width - crop_width) * channels;
  uint8_t* temp = nullptr;
  if(offset > 0) temp = new uint8_t[cinfo.output_width * channels];

  while (cinfo.output_scanline < crop_y + crop_height) {
    /* jpeg_read_scanlines expects an array of pointers to scanlines.
     * Here the array is only one element long, but you could ask for
     * more than one scanline at a time if that's more convenient.
     */
    if(offset>0){
      jpeg_read_scanlines(&cinfo, &temp, 1);
      memcpy(ptr, temp + offset, stride);
    }
    else
      jpeg_read_scanlines(&cinfo, &ptr, 1);
    ptr += stride;
  }
  if(offset > 0){
    delete[] temp;
    temp = nullptr;
  }
  if (cinfo.output_scanline < cinfo.output_height) {
    // Skip the rest of scanlines, required by jpeg_destroy_decompress.
    jpeg_skip_scanlines(&cinfo,
                        cinfo.output_height - crop_y - crop_height);
  }
  jpeg_finish_decompress(&cinfo);
  jpeg_destroy_decompress(&cinfo);
  return tensor.permute({2, 0, 1});
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("decode_and_crop_jpeg",&decode_and_crop_jpeg,"decode_and_crop_jpeg");
}

在下一节中,我们将按照 PyTorch 教程中的步骤将其转换为可在预处理管道中使用的 PyTorch 运算符。

部署 PyTorch 扩展

如 PyTorch 教程中所述,部署自定义运算符有不同的方法。您的部署设计中可能需要考虑许多因素。以下是我们认为重要的一些示例:

  1. 及时编译:为了确保我们的 C++ 扩展是针对我们训练时使用的同一版本的 PyTorch 进行编译的,我们对部署脚本进行了编程,以便在训练环境中进行训练之前编译代码。
  2. 多进程支持:部署脚本必须支持从多个进程(例如,多个 DataLoader 工作线程)加载我们的 C++ 扩展的可能性。
  3. 托管培训支持:由于我们经常在托管培训环境(例如 Amazon SageMaker)中进行培训,因此我们要求部署脚本支持此选项。 (有关定制托管培训环境主题的更多信息,请参阅此处。)

在下面的代码块中,我们定义了一个简单的 setup.py 脚本,用于编译和安装我们的自定义函数,如此处所述。

from setuptools import setup
from torch.utils import cpp_extension

setup(name='decode_and_crop_jpeg',
      ext_modules=[cpp_extension.CppExtension('decode_and_crop_jpeg', 
                                              ['decode_and_crop_jpeg.cpp'], 
                                              libraries=['jpeg'])],
      cmdclass={
   'build_ext': cpp_extension.BuildExtension})

我们将 C++ 文件和 setup.py 脚本放在名为 custom_op 的文件夹中,并定义一个 init.py 以确保安装脚本由单个进程运行一次:

import os
import sys
import subprocess
import shlex
import filelock

p_dir = os.path.dirname(__file__)

with filelock.FileLock(os.path.join(pkg_dir, f".lock")):
  try:
    from custom_op.decode_and_crop_jpeg import decode_and_crop_jpeg
  except ImportError:
    install_cmd = f"{sys.executable} setup.py build_ext --inplace"
    subprocess.run(shlex.split(install_cmd), capture_output=True, cwd=p_dir)
    from custom_op.decode_and_crop_jpeg import decode_and_crop_jpeg

最后,我们修改数据输入管道以使用新创建的自定义函数:

from torchvision.datasets.vision import VisionDataset
input_img_size = [533, 800]
class FakeDataset(VisionDataset):
    def __init__(self, transform):
        super().__init__(root=None, transform=transform)
        size = 10000
        self.img_files = [f'{i}.jpg' for i in range(size)]
        self.targets = np.random.randint(low=0,high=num_classes,
                                        size=(size),dtype=np.uint8).tolist()

    def __getitem__(self, index):
        img_file, target = self.img_files[index], self.targets[index]
        with torch.profiler.record_function('decode_and_crop_jpeg'):
            import random
            from custom_op.decode_and_crop_jpeg import decode_and_crop_jpeg
            with open(img_file, 'rb') as f:
                x = torch.frombuffer(f.read(), dtype=torch.uint8)
            h_offset = random.randint(0, input_img_size[0] - img_size)
            w_offset = random.randint(0, input_img_size[1] - img_size)
            img = decode_and_crop_jpeg(x, h_offset, w_offset, 
                                       img_size, img_size)

        if self.transform is not None:
            img = self.transform(img)
        return img, target

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

transform = T.Compose(
    [RandomMask(),
     ConvertColor(),
     Scale()])

结果

经过我们描述的优化后,我们的步长时间从 0.72 秒降至 0.48 秒,性能提升了 50%!当然,我们优化的影响与原始 JPEG 图像的大小和我们选择的裁剪大小直接相关。

总结

数据预处理管道中的瓶颈很常见,可能会导致 GPU 饥饿并减慢训练速度。考虑到潜在的成本影响,您必须拥有各种工具和技术来分析和解决这些问题。在这篇文章中,我们回顾了通过创建自定义 C++ PyTorch 扩展来优化数据输入管道的选项,展示了其易用性,并展示了其潜在影响。当然,这种优化机制的潜在收益会根据项目和性能瓶颈的细节而有很大差异。

相关文章
|
20天前
|
机器学习/深度学习 数据采集 算法
构建高效图像分类模型:深度学习在处理大规模视觉数据中的应用
随着数字化时代的到来,海量的图像数据被不断产生。深度学习技术因其在处理高维度、非线性和大规模数据集上的卓越性能,已成为图像分类任务的核心方法。本文将详细探讨如何构建一个高效的深度学习模型用于图像分类,包括数据预处理、选择合适的网络架构、训练技巧以及模型优化策略。我们将重点分析卷积神经网络(CNN)在图像识别中的运用,并提出一种改进的训练流程,旨在提升模型的泛化能力和计算效率。通过实验验证,我们的模型能够在保持较低计算成本的同时,达到较高的准确率,为大规模图像数据的自动分类和识别提供了一种有效的解决方案。
|
2月前
|
机器学习/深度学习 人工智能 自动驾驶
深度学习-数据增强与扩充
深度学习-数据增强与扩充
57 1
|
1月前
|
机器学习/深度学习 数据采集 PyTorch
使用PyTorch解决多分类问题:构建、训练和评估深度学习模型
使用PyTorch解决多分类问题:构建、训练和评估深度学习模型
使用PyTorch解决多分类问题:构建、训练和评估深度学习模型
|
2天前
|
机器学习/深度学习 PyTorch API
|
4天前
|
机器学习/深度学习 算法 计算机视觉
深度学习在图像识别中的应用及优化策略
【4月更文挑战第8天】 在计算机视觉领域,深度学习技术已成为推动图像识别进步的关键力量。本文章旨在探讨深度学习模型在图像识别任务中的应用,并分析其性能提升的优化方法。通过对比传统机器学习方法,本文阐述了深度神经网络如何通过多层次特征提取有效识别复杂图像,并讨论了数据增强、网络结构调整、正则化技巧等优化策略。此外,文中还涉及了迁移学习与多任务学习在图像识别中的实际应用案例,以及未来发展趋势。
|
18天前
|
机器学习/深度学习 算法 大数据
基于PyTorch对凸函数采用SGD算法优化实例(附源码)
基于PyTorch对凸函数采用SGD算法优化实例(附源码)
26 3
|
20天前
|
机器学习/深度学习 数据采集 自动驾驶
利用深度学习优化图像识别在自动驾驶系统中的应用
在自动驾驶技术迅猛发展的当下,图像识别作为其核心技术之一,对于提升车辆的环境感知能力至关重要。本文聚焦于探讨如何通过深度学习算法优化图像识别过程,以增强自动驾驶系统的准确性和实时反应能力。文中介绍了卷积神经网络(CNN)在图像处理中的关键作用,分析了数据预处理、模型训练策略以及模型压缩等技术对性能的影响。此外,还探讨了迁移学习在缺乏标注数据时的应用,以及对抗性网络在提高模型鲁棒性方面的潜力。通过实验评估,本文展示了这些技术在真实世界数据集上的应用效果,并对未来自动驾驶系统中图像识别技术的发展趋势进行了展望。
22 1
|
25天前
|
机器学习/深度学习 分布式计算
ICLR 2024:首个零阶优化深度学习框架
【2月更文挑战第28天】ICLR 2024:首个零阶优化深度学习框架
14 1
ICLR 2024:首个零阶优化深度学习框架
|
1月前
|
机器学习/深度学习 算法 网络架构
基于深度学习的图像识别优化策略
【2月更文挑战第21天】 随着人工智能技术的飞速发展,深度学习在图像识别领域取得了突破性进展。然而,在实际应用中,模型的识别效率和准确性常常受限于数据量、计算资源和算法设计。本文旨在探讨针对现有深度学习模型的图像识别优化策略,通过改进训练过程、网络结构与后处理技术,提高模型性能并减少计算资源的消耗。
|
1月前
|
机器学习/深度学习 数据处理 计算机视觉
深度学习在图像识别中的应用及优化策略
【2月更文挑战第18天】 随着计算机视觉技术的迅猛发展,深度学习已成为推动图像识别领域进步的核心力量。本文将探讨深度学习在图像识别任务中的关键应用,并重点分析数据增强、网络结构优化以及迁移学习等提升模型性能的策略。通过深入剖析这些技术,我们旨在为读者提供一套实用的方法论,以应对不断变化的图像识别挑战。
14 3