一、pytorch如何保存訓練模型

最初pytorch用pickle模塊進行臨時保存程序中的對象（比如訓練好的模型、中間計算結果），方便後續直接加載使用。pickle是 Python 的序列化 / 反序列化模塊，核心作用是把 Python 對象（比如字典、列表、類實例等）轉換成字節流（序列化，稱為 “pickling”），或者把字節流恢復成原來的 Python 對象（反序列化，稱為 “unpickling”）。

但是pickle進行模型臨時保存會面臨一些明顯的安全問題，惡意構造的 pickle 數據在反序列化（unpickle）時會執行任意代碼，如下圖所示：

因此，weights_only參數被引入，以限制反序列化的內容，僅允許加載張量（torch.Tensor）、基本數據類型（int/float/str 等）、簡單容器（dict/list/tuple），禁止加載自定義類、函數、複雜對象等可能包含可執行代碼的內容。作為對比，

二、weights_only為什麼有效

這裏我們解釋為什麼weights_only為什麼有效之前，需要先説明為什麼pickle不安全，pickle 裏最危險的兩類能力GLOBAL和REDUCE，GLOBAL可以引用任意模塊裏的任意對象（函數/類），例如 os.system、subprocess.Popen、builtins.eval，REDUCE可以調用一個函數（或可調用對象）並傳入參數來“構造對象”。這個“函數”如果被引用成 os.system 之類，就等於反序列化過程中直接執行命令。
所以普通的 torch.load()（weights_only=False 的老邏輯）在讀取惡意 .pt/.pth 時，會把這些 pickle 指令照做，導致典型的 反序列化 RCE。

而weights_only的有效機制也是主要禁用這兩個函數，如下圖所示，如果pickle指令裏面有['sys','os','posix','nt'] 這種高危模塊黑名單，就會直接raise UnpicklingError。如果利用白名單機制則可以進一步限制危險函數。

三、TorchScript 的風險

weights_only看似杜絕了pytorch在加載模型時存在反序列化漏洞的風險，但是通過分析torch.load()函數，能夠發現這裏有一個分流的邏輯，如果不是 TorchScript zip，才直接輪到 weights_only 起作用。而如果 zip 被識別為 TorchScript（JIT）格式，會先直接執行torch.jit.load()，並不執行weights_only的安全限制。

所以如果我們能夠構建一個TorchScript zip 文件，通過 _is_torchscript_zip() 的格式檢測，但內部包含可被 pickle / Python 反序列化處理的惡意結構，就能實現反序列化漏洞利用。那麼首先我們要説明什麼是TorchScript ，TorchScript 是 PyTorch 為跨語言、無 Python 環境執行而設計的中間表示（IR），其初衷是避免 Python 動態執行與 pickle 風險。如下圖所示

為了構建這樣惡意的TorchScript zip文件，我們首先要了解TorchScript長什麼樣。從原始的python代碼到最後TorchScript的IR圖，大概需要經過四步：

1.python代碼->python AST的轉化：Python 解釋器會先把代碼解析成Python 的抽象語法樹（AST），這是 Python 對代碼結構的 “結構化表示”，這裏能看到函數定義、參數、條件分支、返回語句等內容，但不是pytorch JIT能直接用的格式。

2.python AST->JIT AST(TorchScript 的中間表示):這是 PyTorch 專用的結構化表示，包含了 JIT 能理解的 “函數定義”“參數”“變量類型”“運算邏輯” 等信息,轉換後，JIT 可以對這段代碼做編譯優化、跨平台部署

3.JIT AST->original IR graph:從 “語法結構” 到 “計算執行圖” 的翻譯,核心步驟是遍歷 AST 解析語法 → 生成 IR 節點和數據流 → 處理控制流 → 填充類型信息

4.original IR graph->optimized IR graph:PyTorch JIT 的優化器會對原始 IR 做靜態分析 + 代碼簡化,消除了冗餘的變量加載 / 存儲,簡化了控制流邏輯,保留核心計算邏輯

TorchScript 的序列化流程，即把編譯好的模型（IR、函數）保存為文件，方便後續加載 / 部署的過程。首先是“IR / 函數 → 保存為文件 → 恢復為模型” 的閉環，輸入：已經編譯好的IR graph（中間表示）、ScriptFunction（TorchScript 函數）；封裝：這些內容會被打包成ScriptModule（TorchScript 的模塊對象）；保存：通過torch.save()把ScriptModule存為文件（如module.pt）；加載：之後用torch.load()可以把文件恢復為ScriptModule，直接使用；在執行torch.save()時，會有操作如下：ScriptFunction（函數）添加到ScriptModule中；通過ScriptModuleSerializer::serialize()（序列化器），把ScriptModule拆分為

3 部分存儲：data.pkl：保存模型的值信息（用IValue表示，是 TorchScript 中統一的 “值” 類型）；code/目錄：保存Python 代碼 / 調試信息（方便後續導出可讀代碼）；constants.pkl：保存模型中的常量張量（比如代碼裏的固定參數）。

那麼後續在執行torch.load()操作時，會有重新反序列化的操作，核心邏輯是通過調用ScriptModuleDeserializer::deserialize 這個方法，調用 readArchive 方法進行反序列化，讀取 constants.pkl、data.pkl 文件

在之前TorchScript構建的説明時，從最後的optimized IR graph也能看出，裏面的核心計算邏輯不再是python函數，而是torch.add, torch.mul, aten::add, prim::If 此類Operator（算子），由 C++ / ATen / JIT runtime 實現，在 TorchScript IR 裏以 OP 指令的形式出現，Operators 通過 RegisterOperators 註冊。這意味着所有 TorchScript 可用算子必須 顯式註冊，註冊時綁定到具體 C++ 實現。JIT 在執行 IR 時，只能調用已註冊算子，不能隨意調用 Python 函數。這種通過原子方式調用看似安全，但是細究就能發現存在風險點，本次的漏洞主要出現在aten::save和aten::from_file兩個原子方式上。
aten::save用於在 TorchScript 中保存 tensor / object，而aten::from_file則用於從文件加載數據（tensor / storage）為了調用這兩個方法，我們利用torch.from_file和torch.save兩個函數即可。通過在torchscript進行這兩個函數的調用，我們可以即可實現RCE等漏洞。

五、艾體寶Mend.io價值

從此次漏洞事件來看，AI項目中的許多安全漏洞源於第三方庫，尤其是深度學習框架（如 PyTorch、TensorFlow）和數據處理工具（如 Pandas、NumPy）。Mend.io 可以深入分析項目中的所有第三方依賴，識別其中的已知漏洞和安全隱患。通過自動化漏洞掃描和依賴分析，Mend.io 能幫助開發團隊及時發現並修復潛在的安全風險。