チェックポイントファイルの拡張子について

くろはです。今日は画像生成AI関連でチェックポイント(モデルファイル)の拡張子について疑問に思って調べてみたことを共有したいと思います。 Googleなどで画像生成AIについて検索をするとおそらく目にするであろう「StableDiffusion」、このAIではチェックポイントという学習済データを指定することによって画風などをある程度決めることが可能です。そしてそのチェックポイントファイルは「Civitai」や「Hugging Face」といったウェブサイト上からダウンロードすることが可能です。 実際にCivitaiからチェックポイントファイルをDLしてみたところ、ファイルの拡張子が「.ckpt」と「.safetensors」というように複数あることに気づきました。今回はこの2種類の拡張子について軽くまとめます。

ckpt

「.ckpt」という拡張子は「pickle」というPythonモジュールを用いて直列化して保存されたデータに用いられる、Python固有のデータフォーマットです。pickleモジュールはPythonの各種オブジェクトをバイト列に変換したり、変換したりしたバイト列をファイルに保存したりする用途に使われます。 pickleの主な利点はPythonオブジェクトの状態を保存して後で再利用できることで、機械学習や大規模なデータ処理に非常に役立ちます。そんなpickleですがモジュールのトップレベルで定義されている組み込み関数などがPickle化できるため任意のコードが実行される可能性があり、脆弱性の問題を抱えています。この問題について、Pythonの公式ドキュメントでも注意喚起をしています。

safetensors

一方、「.safetensors」という拡張子は上で名前の出たHugging Faceによって開発された新しいシリアル化形式です。深層学習で重要なテンソルという呼ばれる大きくて複雑なデータの塊を保存、および取得するための特別な方法で、ファイルにはテンソルを安全に処理するためのアルゴリズムが保存されており悪意のあるコードに対して安全であると考えられています。 補足ですがテンソルとは、物理学やコンピューターサイエンスなどの様々な分野で使用される数学的な概念です。データを多次元配列として表現する方法です。ディープラーニングと人工知能の領域においてはテンソルはデータの編成と操作に使用される基本的なデータ構造として見られています。   以上、簡素ではありますが画像生成AIについて調べている過程で出会った2種類の拡張子についての共有でした。