[機械学習のはじめ方] Part47: モデル保存（joblib, pickle, torch.save）

時間のかかる学習は一度だけ！モデル保存の基本をマスターしよう✨

なぜモデルを保存する必要があるの？ 🤔

機械学習モデルの学習には、特に大規模なデータセットを使用する場合、多くの時間と計算リソースが必要になることがあります。せっかく学習させたモデルを、プログラムを実行するたびに再学習するのは非効率ですよね。

そこで登場するのが「モデルの保存」です！学習済みのモデルをファイルとして保存しておくことで、以下のようなメリットがあります。

時間の節約: 再学習の必要がなくなり、すぐにモデルを利用できます。
リソースの節約: 学習にかかる計算コストを削減できます。
再現性の確保: 同じ学習済みモデルを使って、いつでも同じ予測結果を得られます。
デプロイメント: 保存したモデルをWebアプリケーションやAPIに組み込んで、実際のサービスで利用できます。
共有と共同作業: 学習済みモデルを他の人と共有しやすくなります。

今回は、Pythonでよく使われるモデル保存の方法として、pickle、joblib、そしてPyTorchのtorch.saveについて解説していきます。それぞれの特徴を理解し、状況に応じて使い分けられるようになりましょう！

1. pickle: Python標準のシリアライズライブラリ 🥒

pickleは、Pythonに標準で組み込まれているライブラリで、Pythonオブジェクトをバイトストリームに変換（シリアライズ）してファイルに保存したり、逆にファイルから読み込んでオブジェクトを復元（デシリアライズ）したりできます。学習済みモデルもPythonオブジェクトなので、pickleを使って保存・読み込みが可能です。

使い方

使い方は非常にシンプルです。pickle.dump()で保存し、pickle.load()で読み込みます。


import pickle
from sklearn.linear_model import LinearRegression
import numpy as np

# 簡単なモデルを作成して学習
X = np.array([[1], [2], [3]])
y = np.array([2, 4, 6])
model = LinearRegression()
model.fit(X, y)

# モデルをpickle形式で保存 (バイナリ書き込みモード 'wb')
with open('linear_model.pkl', 'wb') as f:
    pickle.dump(model, f)

# 保存したモデルを読み込み (バイナリ読み込みモード 'rb')
with open('linear_model.pkl', 'rb') as f:
    loaded_model = pickle.load(f)

# 読み込んだモデルで予測
print(loaded_model.predict(np.array([[4]])))

メリット

Pythonの標準ライブラリなので、追加のインストールが不要。
モデルだけでなく、リストや辞書など、ほとんどのPythonオブジェクトを保存できる。

デメリット・注意点 ⚠️

セキュリティリスク: pickleは任意のコードを実行できてしまう脆弱性があります。信頼できないソースから提供された.pklファイルを読み込むのは非常に危険です。絶対にやめましょう。
効率性: NumPy配列などの大きなデータを扱う場合、joblibと比較して効率が悪い（速度が遅い、ファイルサイズが大きい）ことがあります。
互換性: Pythonのバージョン間で互換性がない場合があります。また、保存時と異なる環境（ライブラリのバージョンが違うなど）で読み込むとエラーになることがあります。
可読性: 保存されたファイルはバイナリ形式なので、人間が直接中身を読むことは困難です。

Python 3.8以降では、pickleのプロトコル5により、大きなNumPy配列の扱いの効率が改善されていますが、依然としてjoblibの方が推奨される場面が多いです。

2. joblib: Scikit-learn推奨の保存方法 🚀

joblibは、科学技術計算、特に機械学習の分野でよく使われるライブラリです。pickleをベースにしていますが、特にNumPy配列のような大きなデータを効率的に扱えるように最適化されています。Scikit-learnの公式ドキュメントでも、モデルの保存にはjoblibの使用が推奨されています。

使い方

joblibを使うには、まずインストールが必要です（pip install joblib）。使い方はpickleとよく似ており、joblib.dump()で保存し、joblib.load()で読み込みます。


import joblib
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris

# データの準備とモデルの学習
iris = load_iris()
X, y = iris.data, iris.target
model = RandomForestClassifier()
model.fit(X, y)

# モデルをjoblib形式で保存
joblib.dump(model, 'random_forest_model.joblib')

# 保存したモデルを読み込み
loaded_model = joblib.load('random_forest_model.joblib')

# 読み込んだモデルで予測
print(loaded_model.predict(X[:5]))

メリット

効率性: NumPy配列などの大きなデータを含むオブジェクト（Scikit-learnの学習済みモデルなど）の保存・読み込みがpickleよりも高速で、メモリ効率も良いことが多いです。
圧縮オプション: 保存時に圧縮レベルを指定でき、ファイルサイズを小さく抑えることができます。
メモリマッピング: `mmap_mode`オプションを使うと、モデル全体をメモリに読み込まずに、ファイルから直接必要な部分だけを読み込めるため、巨大なモデルを扱う際や、複数のプロセスでモデルを共有する際にメモリ使用量を抑えられます。

デメリット・注意点 ⚠️

外部ライブラリなので、別途インストールが必要です。
pickleと同様に、信頼できないファイルを開くことにはセキュリティリスクが伴います（ただし、pickleよりは安全性が高いとされています）。
バージョン間の互換性の問題はpickleと同様に存在します。

Scikit-learnのモデルを保存する場合は、joblibが第一候補となるでしょう。

3. torch.save: PyTorchモデルの保存 🔥

PyTorchを使ってディープラーニングモデルを開発している場合、torch.save()関数を使ってモデルを保存するのが一般的です。PyTorchには、モデルの保存に関して主に2つのアプローチがあります。

方法1: モデルのパラメータ（state_dict）のみを保存 (推奨)

最も推奨される方法は、モデルの学習可能なパラメータ（重みやバイアスなど）だけをstate_dictと呼ばれるPythonの辞書オブジェクトとして保存する方法です。state_dictには、モデルの各層とそのパラメータのテンソルがマッピングされています。この方法が推奨される理由は、柔軟性が高く、後でモデルの構造を変更した場合でも、互換性のある部分のパラメータを読み込むことができるためです。


import torch
import torch.nn as nn

# 簡単なモデルを定義
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 2)

    def forward(self, x):
        return self.linear(x)