flax.serialization 套件

Jax 序列化工具程式。

所有具有狀態的 Flax 類別（例如，最佳化器）都可以轉換成 numpy 陣列的狀態字典，以便輕鬆地序列化。

狀態字典

flax.serialization.from_state_dict(target, state, name='.')

使用狀態字典復原給定目標的狀態。

此函數將目前的目標用作參數。這讓我們知道目標的確切結構，還能讓我們新增斷言，以避免形狀和資料類型改變。

實際上，不使用target 中任一個樹狀結構中的值。只使用樹狀結構、形狀和資料類型。

參數

• target – 要復原狀態的物件。

• state – 由to_state_dict 生成的字典，其中包含target 所需的新狀態。

• name – 走過的樹枝名稱，用於改善序列化錯誤訊息。

傳回結果

復原狀態之後的物件副本。

flax.serialization.to_state_dict(target)

傳回一個字典，其中包含給定目標的狀態。

flax.serialization.register_serialization_state(ty, ty_to_state_dict, ty_from_state_dict, override=False)

註冊要序列化的類型。

參數

• ty – 要註冊的類型

• ty_to_state_dict – 一個接受 ty 實體，並以字典形式傳回其狀態的函式。

• ty_from_state_dict – 一個接受 ty 實體和狀態字典的函式，並傳回狀態已還原的實體副本。

• override – 覆寫先前註冊的序列化處理常式（預設：False）。

訊息格式化套件序列化

flax.serialization.msgpack_serialize(pytree, in_place=False)

儲存資料結構為訊息格式化套件格式的位元組。

低階函式，僅適用於具有陣列節點的 Python 結構體；若有自訂物件，請使用 to_bytes。其會將大於 MAX_CHUNK_SIZE 的陣列分割為多個區塊。

參數

• pytree – 包含 Python 原生物件和陣列節點的字典、清單、或元組的 Python 結構體。

• in_place – 布林值，說明是否直接修改 pytree。

傳回結果

pytree 的訊息格式化套件編碼位元組。

flax.serialization.msgpack_restore(encoded_pytree)

從訊息格式化套件格式的位元組中還原資料結構。

低階函式，僅適用於具有陣列節點的 Python 結構體；若有自訂物件，請使用 from_bytes。

參數

encoded_pytree – Python 結構體的訊息格式化套件編碼位元組。

傳回結果

包含 Python 原生物件和陣列節點的字典、清單、或元組的 Python 結構體。

flax.serialization.to_bytes(target)

儲存最佳化器或其他物件為訊息格式化套件串列化狀態字典。

參數

target – 含狀態字典註冊項的範本物件，可序列化為訊息格式化套件格式。通常為 flax 模型或最佳化器。

傳回結果

target 物件的訊息格式化套件編碼狀態字典位元組。

flax.serialization.from_bytes(target, encoded_bytes)

從 msgpack 序列化狀態字典還原最佳化器或其他物件。

參數

• target - 具有狀態字典註冊，且與從 encoded_bytes 解序列化的結構相符範本物件。

• encoded_bytes - 與 target 在結構上同構的 msgpack 序列化物件。通常是 flax 模型或最佳化器。

傳回結果

在結構上與 target 同構的新物件，包含儲存資料的更新之葉資料。