PyTorch 2.6 のセキュリティ: モデルチェックポイントロードの安全対策

PyTorch 2.6 のセキュリティ: モデルチェックポイントロードの安全対策

こんにちは!

今日は、魅力的なPyTorchチェックポイントが配られているからと言って無邪気に使っちゃうと、超ヤバイよ、というお話になります。

みなさまモデルチェックポイントをロードする際のセキュリティリスクについて考えたことはありますでしょうか。実はモデルの重みファイルだとばかり思っていたチェックポイントが、思わぬセキュリティホールになる可能性があります。

本記事では、PyTorchのtorch.load関数の安全性と、モデルチェックポイントを適切に扱うための実践的なガイドラインを紹介します。

モデルチェックポイントの隠れた危険性

PyTorchのtorch.load関数は非常に便利な一方で、セキュリティ上の重大なリスクを含んでいます。

その理由は、

  • チェックポイント単なるパラメータだけではないよ!
    チェックポイントファイルには、モデルの重み(weights)だけでなく、クラスや関数など任意のPythonコードを含めることが可能です。
  • 実行可能なコードが入ってるよ!
    これは、チェックポイントが単なる「データファイル」ではなく、Pythonコードを実行できる可能性を持つファイルであることを意味します。
  • ってことで潜在的な脆弱性があるよ!
    悪意あるチェックポイントを不用意にロードすると、ファイルをロードした瞬間に任意のコードが実行される可能性があります。

そうです、チェックポイントは重みデータだけではないんです。

チェックポイントファイルに入れられるもの

ということで、PyTorchのチェックポイントファイル(.ptや.pth,.ckptファイル)には、モデルの重み(パラメータ)だけでなく、実際にPythonのコードも含めることができてしまいます。

これは、PyTorchが内部的にPythonのpickleシリアライゼーション形式を使用しているためです。

pickleは、Pythonオブジェクトの状態を保存・復元するための仕組みですが、クラス定義や関数など、実行可能なコードも含めて保存できる特性があります。

具体的には、チェックポイントには以下のような要素が含まれてる可能性があります。

  • モデルの重みパラメータ
  • モデルのアーキテクチャ情報
  • オプティマイザの状態
  • 学習率スケジューラの状態
  • カスタムクラス・関数の定義(←ここが危ない)
  • その他のメタデータ

「カスタムクラス・関数の定義」が、セキュリティ上の懸念となります。悪意のある人が作成したチェックポイントには、ファイル削除やシステムコマンドの実行など、危険なコードが含まれている可能性があります。

だから、魅力的なチェックポイントが配られているからと言って無邪気に使っちゃうと、超ヤバイんです。

具体的な攻撃シナリオ

悪意あるチェックポイントが引き起こす可能性のある問題は次のようなものがあります

  1. (あなたが)不正なファイルをダウンロードする
  2. (あなたが)それをweights_only=Falseでロードする
  3. (わるい奴が仕込んだ)チェックポイントに含まれた悪意あるコードが即座に実行される
  4. (最悪の場合は)攻撃者はPC上で任意の操作(ファイル削除、情報漏洩、マルウェアのインストールなど)が可能になる

PyTorch 2.6でのセキュリティ強化

PyTorch 2.6以降では、セキュリティ対策が強化されました。具体的な変更点としては、

  • torch.load()weights_only引数がデフォルトでTrueに設定されるようになりました
  • これにより、デフォルトでモデルのパラメータ(重み)のみが安全にロードされ、潜在的に危険なコードの実行が防止されます

この変更により、以下のようなエラーが発生したり

RuntimeError: ('Attempted to deserialize object from torch.nn.Module that contains non-parameter/buffer types, which could potentially lead to security vulnerabilities. ...

 _pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint.
         (1) In PyTorch 2.6, we changed the default value of the weights_only argument in torch.load from False to True. Re-running torch.load with weights_only set to False will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
         (2) Alternatively, to load with weights_only=True please check the recommended steps in the following error message.
         WeightsUnpickler error: Unsupported global:...

このエラーは、チェックポイントファイルに単なるモデルの重みだけでなく、クラスや関数などのコード情報も含まれている場合に発生します。例えば、fairseq.data.dictionary.Dictionaryのようなクラスが含まれていると、デフォルト設定では未許可のクラスとして拒否されます。

安全の基準:何が「安全」で何が「安全でない」か

🟢 「安全」と判断できるファイル

次のような場合はほぼ「安全」と考えることができます

  • 自分自身で作成したモデルのチェックポイント(内容を完全に把握しているため)
  • 信頼できる公式ソースから直接取得したモデルファイル(有名な大学、著名な企業、公式レポジトリなど)
  • 公式のGitHubリポジトリや公式ウェブサイトから取得したモデルで、提供者が広く知られており、モデルのチェックポイントの内容が明示的に示されている場合(例:Hugging Faceの公式リポジトリ、Facebookのfairseq公式リポジトリなど)

(とはいえ、注意深くセキュリティ情報は把握しておく必要はありますし、基本は weights_only =True でロードです。)

🔴 「安全でない」可能性のあるファイル

次のような場合は「安全でない」と疑うべきです:

  • 出所不明のサイトや第三者が匿名で共有したモデル
  • SNSや非公式フォーラムなど、素性がわからない人物から直接ダウンロードしたモデルファイル
  • モデルファイルの配布者が特定できず、かつモデルの構造やコードについて十分な情報が公開されていない場合

このようなファイルは絶対にweights_only=Falseでロードしてはいけません。

🔐 安全なモデルロードのためのコード例と対処法

エラーが発生した場合の解決方法

ここではfacebookの fairseqを使っていて実際に発生したエラーと対処法についてみてみます

 _pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint.
         (1) In PyTorch 2.6, we changed the default value of the weights_only argument in torch.load from False to True. Re-running torch.load with weights_only set to False will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
         (2) Alternatively, to load with weights_only=True please check the recommended steps in the following error message.
         WeightsUnpickler error: Unsupported global: GLOBAL fairseq.data.dictionary.Dictionary was not an allowed global by default. Please use torch.serialization.add_safe_globals([Dictionary]) or the torch.serialization.safe_globals([Dictionary]) context manager to allowlist this global if you trust this class/function.

このエラーは、PyTorchのtorch.load()でモデルのチェックポイントを読み込む際、セキュリティのためデフォルト設定が厳しくなったために発生したものです。

具体的には、以下の原因によります

  • PyTorch 2.6以降で、torch.load()weights_only引数がデフォルトでTrueとなり、モデルのパラメータ(重み)のみを安全にロードしようとします。
  • このエラーが発生するチェックポイントファイルには、単なるモデルの重みだけでなく、クラスや関数などのコード情報も含まれています。そのため、weights_only=Trueのままロードしようとすると、未許可のクラスや関数が含まれているとして拒否されます。
  • 今回のエラーの場合、未許可として拒否されたクラスは fairseq.data.dictionary.Dictionary となっていることがこのメッセージから読み取れますね。

① 推奨の解決方法(安全性高)

信頼できるクラスを「許可リスト」に登録してから読み込みます。

import torch
from fairseq.data.dictionary import Dictionary

# このクラスが安全と分かっている場合のみ実行
torch.serialization.add_safe_globals([Dictionary])

# その後にチェックポイントをロード
checkpoint = torch.load('checkpoint.pt')

もしくは、一時的に許可する場合はコンテキストマネージャを利用します。

import torch
from fairseq.data.dictionary import Dictionary

with torch.serialization.safe_globals([Dictionary]):
    checkpoint = torch.load('checkpoint.pt')

② 旧来の方法(安全性低)

weights_only=Falseを指定してロードします。ただし、この方法は安全ではないため、信頼できるソースから取得したファイルにのみ使用してください。

import torch

checkpoint = torch.load('checkpoint.pt', weights_only=False)

実装例:安全にチェックポイントをロードする関数

以下は、安全性を考慮したチェックポイントロード関数の例です

import torch
from typing import List, Type, Optional

def load_checkpoint_safely(
    checkpoint_path: str,
    trusted_classes: Optional[List[Type]] = None,
    force_weights_only: bool = False
) -> dict:
    """
    安全にチェックポイントをロードする関数
    
    Args:
        checkpoint_path: チェックポイントファイルのパス
        trusted_classes: 信頼できるクラスのリスト(None の場合は空リスト)
        force_weights_only: True の場合、weights_only=True で強制的にロード
        
    Returns:
        チェックポイントの内容
        
    Raises:
        RuntimeError: ロードに失敗した場合
    """
    if force_weights_only:
        return torch.load(checkpoint_path, weights_only=True)
        
    trusted_classes = trusted_classes or []
    
    try:
        # まず安全にロードを試みる
        return torch.load(checkpoint_path)
    except RuntimeError as e:
        if "non-parameter/buffer types" in str(e):
            # 信頼できるクラスがある場合はそれらを使用
            if trusted_classes:
                with torch.serialization.safe_globals(trusted_classes):
                    return torch.load(checkpoint_path)
            else:
                raise RuntimeError(
                    f"チェックポイントに未許可のクラスが含まれています。"
                    f"信頼できるソースの場合は、trusted_classes パラメータを使用してください。"
                    f"原エラー: {str(e)}"
                )
        else:
            # その他のエラーはそのまま再発生
            raise

さて、最後に今回の事例についてベストプラクティスをまとめておきましょう

ベストプラクティスのまとめ

  1. デフォルトは安全第一
    • PyTorch 2.6以降のデフォルト設定(weights_only=True)を活用する
    • 信頼性が不明なファイルは常にデフォルト設定のままロードを試みる
  2. 明示的な許可リストの使用
    • 信頼できるクラスはadd_safe_globalsまたはsafe_globalsコンテキストマネージャで明示的に許可する
    • これにより、必要最小限のクラスのみを許可し、安全性を確保できる
  3. weights_only=Falseは最終手段です
    • 信頼できるファイルの場合のみ、最終手段として使用する
    • 不明なソースのチェックポイントには絶対に使用しない

まとめ:セキュリティ意識の向上を

PyTorch 2.6以降のセキュリティ強化は、深層学習モデルの安全な共有と利用を前進させました。

基本的な安全の基準は非常にシンプルで

  • 出所が明確で信頼(安全のエビデンスがある)できる → 安全
  • 出所不明、匿名、不審な提供元 → 安全でない

当社でもプロダクトのベースを PyTorch 2.6 にしたら、いままで動いていたコードが動かなくなり、多くの見直しも発生しました。が、そのおかげで、モデル・セキュリティの意識と知識を一歩深めることができたと感じております。

ディープラーニングの適用例やプレイヤーが増えるにつれ多くの有用なモデルが配布されるようになりましたが、プレイヤーが増えると、このようによからぬことを考えるプレイヤーも増えるという事なので、今後もチェックポイントのロードには細心の注意を払っていきたいと思います。旧バージョンのPyTorchを使ってるプロダクト開発する必要がある場合は特にきをつけるべきですね。

それでは、また次回お会いしましょう!

Read more

OpenCV cv2.imwrite で発生する「_img.empty()」エラーと「動画安定化」による解決法

OpenCV cv2.imwrite で発生する「_img.empty()」エラーと「動画安定化」による解決法

こんにちは! 画像処理や動画解析の現場で広く利用されている OpenCV。 しかし実務で動画処理を行っていると、時折以下のようなエラーに遭遇することがあります。 cv2.error: OpenCV(4.11.0) /io/opencv/modules/imgcodecs/src/loadsave.cpp:929: error: (-215:Assertion failed) !_img.empty() in function 'imwrite' このエラーは、cv2.imwrite() に渡された画像が空(None またはサイズ0) の場合に発生します。 一見単純に見える問題ですが、背後には「入力動画の不安定さ」や「並列処理の競合」といった要因が潜んでいることが少なくありません。 本記事では、このエラーの発生原因を掘り下げ、実務で効果のある解決策として 「動画の安定化(正規化)」 を紹介します。 TL;

By Qualiteg プロダクト開発部
発話音声からリアルなリップシンクを生成する技術 第5回(前編):Transformerの実装と実践的な技術選択

発話音声からリアルなリップシンクを生成する技術 第5回(前編):Transformerの実装と実践的な技術選択

こんにちは!リップシンク技術シリーズもいよいよ終盤となりました。 前回(第4回)では、LSTMの学習プロセスと限界について詳しく解説しました。限られたデータでも効果的に学習できるLSTMの強みを理解する一方で、長距離依存の処理に限界があることも明らかになりました。そして、この問題を解決する革新的なアプローチとして、すべての位置の情報を同時に参照できるTransformerのSelf-Attention機構を紹介しました。 第5回の今回は、 Transformerの具体的なネットワーク設計から始め、その実装上の課題を明らかにします。(前編※) そして、LSTMとTransformerの長所を組み合わせたハイブリッドアプローチを紹介し、実際の製品開発における技術選択の指針を示します。最後に、感情表現への拡張という次なる挑戦についても触れていきます。(後編※) ※Transformerの仕組みは複雑であるため、第5回は前編と後編に分けて解説させていただく予定です。 1. Transformerベースのネットワーク設計 1.1 全体アーキテクチャ図 では、さっそく、Tran

By Qualiteg 研究部, Qualiteg コンサルティング
大企業のAIセキュリティを支える基盤技術 - 今こそ理解するActive Directory 第2回 ドメイン環境の構築

大企業のAIセキュリティを支える基盤技術 - 今こそ理解するActive Directory 第2回 ドメイン環境の構築

こんにちは、今回はシリーズ第2回ドメイン環境の構築 - 検証環境の構築手順について解説いたします! 連載の構成 第1章:基本概念の理解 - Active DirectoryとKerberos/NTLM認証の基礎 【★今回です★】第2章:ドメイン環境の構築 - 検証環境の構築手順 第3章:クライアントとサーバーのドメイン参加 - ドメイン参加の詳細手順 第4章:プロキシサーバーと統合Windows認証 第5章:ブラウザ設定と認証 - 各ブラウザでの設定方法 第6章:トラブルシューティング - よくある問題と解決方法 第7章:セキュリティとベストプラクティス - 本番環境での考慮事項 第8章:実践的な構成例 - AIセキュリティツールとの統合事例 第2章:ドメイン環境の構築 2.1 ドメイン名の設計 2.1.1 ドメイン名の命名規則 Active Directoryを構築する際、

By Qualiteg コンサルティング
AIがよく間違える「クロージャ問題」の本質と対策

AIがよく間違える「クロージャ問題」の本質と対策

こんにちは! 本日は「クロージャ問題」に関する話題となります。 Pythonでループ内に関数を定義したことはありますか? もしあるなら、あれれ?な挙動に遭遇したことがあるかもしれません。 本稿では、Pythonプログラマーなら一度は経験する「クロージャ問題」について、初心者にもわかりやすく解説してみたいとおもいます クロージャとは何か? そもそも ”クロージャ” とは何でしょうか。 クロージャ(closure)とは、関数が自分の定義されたスコープの変数を覚えて持ち運ぶ仕組み のことです。 もう少し分解すると、次の2つがポイントとなります 1. 内側の関数が、外側の関数の変数を使える 2. 外側の関数が終了しても、その変数は生き続ける 普通の関数とクロージャ―を使った関数を比較してみましょう 普通の関数との比較 まずは普通の関数から、 def add(x, y): return x + y print(add(3, 5)) # 8 print(add(3, 7)

By Qualiteg プロダクト開発部