import pandas as pd
import xgboost as xgb
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# データを読み込む
iris = load_iris()
df = pd.DataFrame(data=iris.data, columns=iris.feature_names)
df['class'] = iris.target
display(df)

X = df.drop(columns=['class'])
y = df['class']

# データをトレーニングセットとテストセットに分割
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# モデルのパラメータを設定
params = {
    'objective': 'multi:softmax',
    'num_class': 3,
    'eval_metric': 'merror'
}

# XGBoostのモデル構築
dtrain = xgb.DMatrix(X_train, label=y_train)
num_rounds = 100
model = xgb.train(params, dtrain, num_rounds)

# モデルの保存
pd.to_pickle(model, 'model.pkl')

もちろん、このままの状態では model.pkl を持ち出して使うこともできるし、pd.read_pickleなどで使用している特徴量も確認できてしまう。

2. モデルの保護

pickleファイル　→　base64でテキスト化　→　テキストを暗号化

# バイナリ形式のデータを読み込む
with open("model.pkl", "rb") as file:
    read_model = file.read()
print('<バイナリの状態>\r\n', read_model)

import base64

# PickleデータをBase64でエンコードする
encoded_data = base64.b64encode(read_model)
print('<base64でエンコードした状態>\r\n', encoded_data)

from cryptography.fernet import Fernet

# 鍵を作成して、base64でエンコードしたものを暗号化して txt形式で出力
passphrase = b'\xc8P\xde@\xf8\x0e\xfe\x81\xf5\xdf\x1b\xafLDe\xdf\xdfT\xd5\xab6\xa1\x9f\xf0\xd49\xb8\xce\x92A\x1c\x85'
key = base64.urlsafe_b64encode(passphrase)

f = Fernet(key)
encrypted_data = f.encrypt(encoded_data)
print('<暗号化した状態>\r\n', encrypted_data)

# 暗号化したモデル情報を txt形式で出力
with open("encrypted_model.txt", "wb") as file:
    file.write(encrypted_data)

上記コードのパスフレーズの設定と、keyの作成で超ハマりました。
っというのも、ChatGPTやググったりして出てくるサンプルコードは、key=Fernet.generate_key() の書き方しかなく、一旦ランダムで作った鍵をローカルに保持する形になってしまう。
そうなると、結局クライアントPC上にそのまま使える鍵(decode可能な鍵)が存在することになるので、鍵の在りかさえ分かれば復号ができてしまう。

鍵自体はローカルに置かず、モデル利用のタイミングで同じ鍵を作成して復号化する仕組みを実現したかったため、今回は上記の様なコードにしている。
（上記の様なコードを書くために、cryptographyのパッケージを直に見てみたが、generate_key() のメソッドは下図の様な処理だったので、入れる引数をrandomではなく固定化させてみた）

3. 保護したモデルの使い方イメージ

# 暗号化したモデル情報の txtを読み込む
with open("encrypted_model.txt", "rb") as file:
    encrypted_data = file.read()

# 同じ鍵を作成して復号化
import base64
from cryptography.fernet import Fernet

# 鍵を作成して、base64でエンコードしたものを暗号化して txt形式で出力
passphrase = b'\xc8P\xde@\xf8\x0e\xfe\x81\xf5\xdf\x1b\xafLDe\xdf\xdfT\xd5\xab6\xa1\x9f\xf0\xd49\xb8\xce\x92A\x1c\x85'
key = base64.urlsafe_b64encode(passphrase)

f = Fernet(key)
encoded_data = f.decrypt(encrypted_data)

# PickleデータをBase64でエンコードする
model_info = base64.b64decode(encoded_data)
model_info

# モデルの読み込み
import pickle
model = pickle.loads(model_info)

モデル利用のツールやアプリについて、上記のコードを.pyc形式のコンパイルされた形で提供すれば、最低限の資産流用の防止にはなるのではないか。。。

はじめに

実行環境

今回やりたい事のイメージ

1. モデル構築とpickle保存

2. モデルの保護

3. 保護したモデルの使い方イメージ

Discussion