「カスタムPyTorchオペレーターを使用してDLデータ入力パイプラインを最適化する方法」
Optimizing DL data input pipelines using custom PyTorch operators
PyTorchモデルのパフォーマンス分析と最適化-パート5
この投稿は、GPUベースのPyTorchワークロードのパフォーマンス分析と最適化に関する一連の投稿の5番目です。前回の投稿では、PyTorchプロファイラとTensorBoardを使用して、DLトレーニングワークロードのデータ前処理パイプラインのパフォーマンスボトルネックを特定、分析、解決する方法を示しました。この投稿では、PyTorchのカスタムオペレータの作成に対するサポートについて説明し、それがデータ入力パイプラインのパフォーマンスボトルネックの解決やDLワークロードの高速化、トレーニングコストの削減にどのように役立つかを実証します。この投稿にはYitzhak Levi氏とGilad Wasserman氏への感謝の意があります。この投稿に関連するコードは、このGitHubリポジトリで見つけることができます。
PyTorch拡張機能の構築
PyTorchには、torch.nnをカスタムモジュールや関数で拡張するためのいくつかの方法があります。この投稿では、PyTorchがカスタムC++コードの統合をサポートしていることに興味があります。この機能は重要です。なぜなら、一部の操作はPythonよりもC++で（はるかに）効率的に実装できる場合があるからです。CppExtensionなどの指定されたPyTorchユーティリティを使用すると、これらの操作をPyTorchに「拡張」として簡単に組み込むことができます。PyTorchのコードベース全体を取得して再コンパイルする必要はありません。この機能の背後にある動機と使用方法の詳細については、公式のPyTorchチュートリアル「カスタムC++およびCUDA拡張機能」を参照してください。この投稿では、CPUベースのデータ前処理パイプラインの高速化に興味があるため、C++拡張機能だけで十分です。将来の投稿では、この機能を使用してGPU上で実行されるトレーニングコードの高速化のためのカスタムCUDA拡張機能の実装方法を示すことを希望しています。
おもちゃの例
前回の投稿では、533×800のJPEG画像のデコードから始まり、ランダムな256×256の切り抜きを抽出し、さらにいくつかの変換を行ってトレーニングループに供給するデータ入力パイプラインを定義しました。PyTorchプロファイラとTensorBoardを使用して、ファイルからの画像の読み込みに関連する時間を測定し、デコードの無駄さについて認識しました。完全性のために、以下にコードをコピーします。
import numpy as npfrom PIL import Imagefrom torchvision.datasets.vision import VisionDatasetinput_img_size = [533, 800]img_size = 256class 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()])前回の投稿では、最適化された平均ステップ時間は0.72秒であることがわかりました。おそらく、興味のある切り抜きのみをデコードできれば、パイプラインはより速く実行されるでしょう。残念ながら、この投稿の執筆時点では、PyTorchにはこの機能をサポートする関数が含まれていません。しかし、カスタムオペレーションの作成ツールを使用することで、独自の関数を定義して実装することができます！
- これらの4つのパッケージで、あなたの探索的データ分析を簡素化する
- 「過去のデータ、Ray、およびAmazon SageMakerを使用して装置のパフォーマンスを最適化する」
- 「PythonとMatplotlibを使用して米国のデータマップを作成する方法」
カスタムJPEG画像デコードおよび切り抜き関数
libjpeg-turboライブラリは、libjpegと比較してさまざまな強化と最適化を含むJPEG画像コーデックです。特に、libjpeg-turboには、jpeg_skip_scanlinesとjpeg_crop_scanlineなどの関数が含まれており、画像内の事前定義された切り抜きのみをデコードすることができます。conda環境で実行している場合は、次のコマンドを使用してインストールできます。
conda install -c conda-forge libjpeg-turbo注意：libjpeg-turboは、以下の実験で使用する公式のAWS PyTorch 2.0 Deep Learning 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>(); // デコンプレッションの構造体をセットアップ cinfo.err = jpeg_std_error(&jerr.pub); jerr.pub.error_exit = torch_jpeg_error_exit; /* my_error_exitが使用するsetjmpの戻りコンテキストを確立します。 */ setjmp(jerr.setjmp_buffer); jpeg_create_decompress(&cinfo); torch_jpeg_set_source_mgr(&cinfo, datap, data.numel()); // ヘッダから情報を読み取ります。 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) { // 残りのスキャンラインをスキップします。これは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");}完全なC++ファイルはこちらで入手できます。
次のセクションでは、PyTorchチュートリアルの手順に従い、このコードをPyTorchオペレータに変換し、前処理パイプラインで使用できるようにします。
PyTorch拡張のデプロイ
PyTorchチュートリアルで説明されているように、カスタムオペレータを展開する方法はさまざまです。デプロイ設計に影響を与える要素もいくつかあります。以下に重要だと考えるいくつかの例を示します。
- Just in timeコンパイル：C++拡張がトレーニングに使用するPyTorchのバージョンと同じバージョンでコンパイルされるようにするため、デプロイスクリプトをトレーニング環境内でトレーニングの直前にコンパイルするようにプログラムします。
- マルチプロセスサポート：デプロイスクリプトは、C++拡張が複数のプロセス（たとえば、複数のDataLoaderワーカー）からロードされる可能性があることをサポートする必要があります。
- マネージドトレーニングサポート：通常、Amazon SageMakerなどのマネージドトレーニング環境でトレーニングを行うため、デプロイスクリプトがこのオプションをサポートする必要があります。（マネージドトレーニング環境のカスタマイズについては、こちらを参照してください。）
以下のコードブロックで、ここで説明されているように、カスタム関数をコンパイルしてインストールする簡単なsetup.pyスクリプトを定義します。
from setuptools import setupfrom torch.utils import cpp_extensionsetup(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})カスタム_opという名前のフォルダにC++ファイルとsetup.pyスクリプトを配置し、__init__.pyを定義して、セットアップスクリプトが一度だけ実行され、一つのプロセスで実行されるようにします。
import osimport sysimport subprocessimport shleximport filelockp_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 VisionDatasetinput_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拡張を作成してデータ入力パイプラインを最適化するオプションを検討し、その使いやすさと潜在的な影響を示しました。もちろん、この種の最適化メカニズムから得られる潜在的な利益は、プロジェクトとパフォーマンスボトルネックの詳細に応じて大幅に異なります。
次は何をすればよいですか？ここで議論した最適化技術は、ブログ記事で多く取り上げてきた入力パイプラインの最適化方法の一部です。ぜひチェックしてみてください（たとえば、ここから始めてください）。
We will continue to update VoAGI; if you have any questions or suggestions, please contact us!
Was this article helpful?
93 out of 132 found this helpful
Related articles
