Zenn
📊

PyTorchのテンソルをシリアライズする3つの方法とベンチマーク分析

2025/02/23に公開
Python
PyTorch
benchmark
numpy
pickle
tech

はじめに

PyTorchで開発をしていると、テンソルデータをバイト列に変換（シリアライズ）して保存や転送を行う必要に迫られることがあります。本記事では、一般的な3つのシリアライズ方法について、その特徴とパフォーマンスを解説します。

シリアライズの3つの方法

1. numpy().tobytes() 

def tensor_to_buffer_numpy(tensor: torch.Tensor) -> bytes:
    if tensor.device.type != "cpu":
        tensor = tensor.cpu()
    return tensor.numpy().tobytes()

def buffer_to_tensor_numpy(buffer: bytes, shape: tuple[int, ...], dtype: np.dtype) -> torch.Tensor:
    array = np.frombuffer(buffer, dtype=dtype).copy()
    return torch.from_numpy(array.reshape(shape))
  • NumPy配列を経由してバイト列に変換する方法
  • メモリコピーが発生するが、NumPyとの相互運用性が高い
  • シンプルで理解しやすい実装

2. data_ptr() 

def tensor_to_buffer_ptr(tensor: torch.Tensor) -> bytes:
    if tensor.device.type != "cpu":
        tensor = tensor.cpu()
    nbytes = tensor.nelement() * tensor.element_size()
    ptr = tensor.data_ptr()
    return ctypes.string_at(ptr, nbytes)

def buffer_to_tensor_ptr(buffer: bytes, shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor:
    return torch.frombuffer(bytearray(buffer), dtype=dtype).reshape(shape)
  • テンソルの生のメモリポインタを直接アクセスする方法
  • 最小限のメモリコピーで高速な変換が可能
  • 低レベルな操作のため、注意深い実装が必要

3. torch.save 

def tensor_to_buffer_save(tensor: torch.Tensor) -> bytes:
    buffer = BytesIO()
    torch.save(tensor, buffer)
    return buffer.getvalue()

def buffer_to_tensor_save(buffer: bytes) -> torch.Tensor:
    return torch.load(BytesIO(buffer))
  • PyTorch標準のシリアライズ機能を使用
  • メタデータ（dtype、device等）も含めて保存可能
  • Pickle形式での保存となるため、セキュリティに注意が必要

ベンチマーク結果と考察

テンソルサイズ別シリアライズ性能（MB/s）

テンソルサイズ numpy().tobytes() data_ptr() torch.save
100x1 724.3 1,305.5 18.8
100x100 42,424.0 59,640.4 1,315.9
5000x5000 20,446.0 20,163.4 3,865.6

小規模テンソル（100x1）での処理では、data_ptr()メソッドが約1.3GB/sと最も高速で、numpy().tobytes()の約724MB/sを大きく上回っています。
一方、torch.saveは約19MB/sと著しく低速です。
しかし、テンソルサイズが大きくなると（5000x5000）、numpy().tobytes()とdata_ptr()はともに約20GB/sとほぼ同等の性能を示し、torch.saveも約4GB/sまで性能が向上します。

パフォーマンス比較

分析から得られた主な知見は以下の通りです：

  1. テンソルサイズによる性能特性

    • 小規模テンソル（100要素以下）ではdata_ptr()が最も高速で、torch.saveは著しく遅い
    • 大規模テンソル（1000x1000以上）ではnumpy().tobytes()とdata_ptr()の性能差が縮小
    • メモリ使用量はnumpy().tobytes()とdata_ptr()が同等で、torch.saveは小規模テンソルでオーバーヘッドが大きい

  2. 実用的な選択基準

    • 単純なデータ転送の場合はdata_ptr()が最適
    • NumPyとの相互運用が必要な場合はnumpy().tobytes()
    • メタデータの保存が重要な場合はtorch.save

  3. 考慮すべき注意点

    • すべての方法でCPUテンソルへの変換が必要
    • data_ptr()は低レベル操作のため、メモリ管理に注意が必要
    • torch.saveはPickleを使用するため、信頼できないソースからのデータ読み込みには注意

まとめ

本記事では、PyTorchテンソルのシリアライズについて3つの方法を比較しました。実際の使用時には、以下の選択基準を推奨します：

  • 高速な単純転送が必要な場合：data_ptr()
  • NumPyとの相互運用性が重要な場合：numpy().tobytes()
  • メタデータの保存が必要な場合：torch.save

参考文献

  1. PyTorchのTensorについて
  2. NumPyのtobytesについて
  3. Python ctypes

ソースコードとraw result

完全なソースコードとベンチマーク結果はこちらです。

Discussion