「Safetensors形式のファイルをPythonで処理したい」
「PyTorchでモデルファイルをもっと速く読み込みたい」
このような場合には、safetensorsがオススメです。
この記事では、テンソルを安全・高速に処理できるsafetensorsについて解説しています。
本記事の内容
- safetensorsとは？
- safetensorsのシステム要件
- safetensorsのインストール
- safetensorsの動作確認
それでは、上記に沿って解説していきます。
safetensorsとは？
まず、Safetensorsの説明が必要です。
Safetensorsは、Tensor（テンソル）を安全・高速に処理するためのフォーマットのことです。
例えば、Hugging Face上で公開されているMidjourney v4 Diffusionを見てみましょう。
公開されているファイルを見ると、次のようなファイルを確認できます。
ckptファイルは、pickle（モジュール）で保存したデータです。
Pickle形式はPythonでデータ保存を行う場合に、よく利用されてきました。
このPickle形式に置き換わるモノとして、Safetensors形式が出てきました。
Safetensors形式は、主にTensor（テンソル）を扱います。
Safetensors形式はPickle形式よりも、安全かつ高速にTensorを処理します。
ここで、やっとsafetensorsの説明です。
safetensorsは、TensorをSafetensors形式として処理するためのライブラリとなります。
safetensorsの開発は、Hugging Faceが行っています。
以上、safetensorsについて説明しました。
次は、safetensorsのシステム要件を説明します。
safetensorsのシステム要件
現時点（2022年12月中旬）でのsafetensorsの最新バージョンは、0.2.5となります。
この最新バージョンは、2022年11月24日にリリースされています。
サポートOSに関しては、以下を含むクロスプラットフォーム対応です。
- Windows
- macOS
- Linux
そして、サポート対象となるPythonのバージョンは以下。
- Python 3.7
- Python 3.8
- Python 3.9
- Python 3.10
また、現状のPython公式開発サイクルは以下となっています。
| バージョン | リリース日 | サポート期限 |
| 3.7 | 2018年6月27日 | 2023年6月27日 |
| 3.8 | 2019年10月14日 | 2024年10月 |
| 3.9 | 2020年10月5日 | 2025年10月 |
| 3.10 | 2021年10月4日 | 2026年10月 |
| 3.11 | 2022年10月25日 | 2027年10月 |
このPython公式開発サイクルに従うなら、Python 3.11までサポートすべきなのでしょう。
でも、 safetensorsはPython 3.11でも問題なく動くとは思います。
safetensorsのシステム要件をまとめます。
Python 3.7以降であることを注意するぐらいです。
以上、safetensorsのシステム要件を説明しました。
次は、safetensorsのインストールを説明します。
safetensorsのインストール
検証は、次のバージョンのPythonで行います。
> python -V Python 3.10.4
まずは、現状のインストール済みパッケージを確認しておきます。
> pip list Package Version ----------- ------- pip 22.3.1 setuptools 65.6.3 wheel 0.37.1
次にするべきことは、pipとsetuptoolsの更新です。
pipコマンドを使う場合、常に以下のコマンドを実行しておきましょう。
python -m pip install --upgrade pip setuptools
では、safetensorsのインストールです。
safetensorsのインストールは、以下のコマンドとなります。
pip install safetensors
safetensorsのインストールは、一瞬で終わります。
どんなパッケージがインストールされたのかを確認します。
> pip list Package Version ----------- ------- pip 22.3.1 safetensors 0.2.5 setuptools 65.6.3 wheel 0.37.1
safetensorsは、依存するパッケージが一つもありません。
そのため、既存環境にも容易にsafetensorsを導入できますね。
以上、safetensorsのインストールを説明しました。
次は、safetensorsの動作確認を説明します。
safetensorsの動作確認
safetensorsの動作確認は、以下の項目毎に説明します。
- 事前準備（PyTorchのインストール）
- Safetensors形式データの保存
- Safetensors形式データの読み込み
それぞれを以下で説明します。
事前準備（PyTorchのインストール）
safetensorsの動作確認には、PyTorchが必要となります。
safetensorsは、Tensor（テンソル）型のデータを高速かつ安全に処理します。
そのため、safetensorsはTensor型のデータありきです。
よって、PyTorchのインストールが必須と言えます。
GPU搭載のマシンであれば、GPU版PyTorchをインストールしましょう。
Safetensors形式データの保存
Safetensors形式データのモデルを保存します。
そのためのコードは、以下。
save.py
import torch
from safetensors.torch import save_file
tensors = {
"embedding": torch.zeros((2, 2)),
"attention": torch.zeros((2, 3))
}
save_file(tensors, "model.safetensors")
上記コードを実行すると、model.safetensorsが作成されます。
Safetensors形式データの読み込み
Safetensors形式データのモデルを読み込みます。
そのためのコードは、以下。
load.py
from safetensors import safe_open
tensors = {}
with safe_open("model.safetensors", framework="pt", device=0) as f:
for k in f.keys():
tensors[k] = f.get_tensor(k)
print(tensors[k])
上記コードを実行すると、コンソールに次のように表示されます。
tensor([[0., 0., 0.],
[0., 0., 0.]], device='cuda:0')
tensor([[0., 0.],
[0., 0.]], device='cuda:0')
「cuda:0」と表示されるのは、GPU対応のPytorchを利用しているからです。
CPU版であれば、「cpu」となるのでしょう。
以上、safetensorsの動作確認を説明しました。
