Safetensors 文件

Torch API

您目前正在檢視 main 版本,這需要從原始碼安裝. 如果您希望使用一般的 pip 安裝,請查看最新的穩定版本 (v0.5.0-rc.0)。
Hugging Face's logo
加入 Hugging Face 社群

並獲得增強的文件體驗

開始使用

Torch API

safetensors.torch.load_file

< >

( filename: typing.Union[str, os.PathLike] device: typing.Union[str, int] = 'cpu' ) Dict[str, torch.Tensor]

參數

  • filename (stros.PathLike) — 包含張量的檔案名稱。
  • device (Union[str, int], 選填, 預設為 cpu) — 加載後張量所需存放的裝置。可用的選項包含所有常見的 Torch 裝置位置。

返回

Dict[str, torch.Tensor]

包含名稱作為鍵 (key),值為 torch.Tensor 的字典

將 safetensors 檔案加載為 torch 格式。

範例

from safetensors.torch import load_file

file_path = "./my_folder/bert.safetensors"
loaded = load_file(file_path)

safetensors.torch.load

< >

( data: bytes ) Dict[str, torch.Tensor]

參數

  • data (bytes) — safetensors 檔案的內容

返回

Dict[str, torch.Tensor]

包含名稱作為鍵 (key),值為 CPU 上 torch.Tensor 的字典

從原始字節 (pure bytes) 將 safetensors 檔案加載為 torch 格式。

範例

from safetensors.torch import load

file_path = "./my_folder/bert.safetensors"
with open(file_path, "rb") as f:
    data = f.read()

loaded = load(data)

safetensors.torch.save_file

< >

( tensors: typing.Dict[str, torch.Tensor] filename: typing.Union[str, os.PathLike] metadata: typing.Optional[typing.Dict[str, str]] = None ) None

參數

  • tensors (Dict[str, torch.Tensor]) — 要保存的張量。張量必須是連續的 (contiguous) 且密集的 (dense)。
  • filename (stros.PathLike) — 我們要保存到的檔案名稱。
  • metadata (Dict[str, str], 選填, 預設為 None) — 您可能想在標頭中保存的僅文字元數據。例如,這對於指定關於底層張量的更多資訊很有用。這純粹是資訊性的,不會影響張量的加載。

返回

None

將張量字典以 safetensors 格式保存到 filename。目前沒有機制可以防止調用者在保存檔案時修改數據,因此在調用 save_file 並同時修改 tensors 字典中引用的張量時請務必小心;這可能會導致檔案損壞。

範例

from safetensors.torch import save_file
import torch

tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))}
save_file(tensors, "model.safetensors")

safetensors.torch.save

< >

( tensors: typing.Dict[str, torch.Tensor] metadata: typing.Optional[typing.Dict[str, str]] = None ) bytes

參數

  • tensors (Dict[str, torch.Tensor]) — 要保存的張量。張量必須是連續的 (contiguous) 且密集的 (dense)。
  • metadata (Dict[str, str], 選填, 預設為 None) — 您可能想在標頭中保存的僅文字元數據。例如,這對於指定關於底層張量的更多資訊很有用。這純粹是資訊性的,不會影響張量的加載。

返回

bytes

代表該格式的原始位元組。

將張量字典保存為 safetensors 格式的原始位元組。

範例

from safetensors.torch import save
import torch

tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))}
byte_data = save(tensors)

safetensors.torch.load_model

< >

( model: Module filename: typing.Union[str, os.PathLike] strict: bool = True device: typing.Union[str, int] = 'cpu' ) `(missing, unexpected)

參數

  • model (torch.nn.Module) — 要加載到的模型。
  • filename (stros.PathLike) — 用於加載檔案的檔案名稱位置。
  • strict (bool, 選填, 預設為 True) — 是否在缺少鍵或出現意外鍵時報錯失敗。當為 false 時,該函數將簡單地返回缺少和意外的名稱。
  • device (Union[str, int], 選填, 預設為 cpu) — 加載後張量所需存放的裝置。可用的選項包含所有常見的 Torch 裝置位置。

返回

`(missing, unexpected)

(List[str], List[str]) missing 是模型中在加載期間未被修改的名稱;unexpected 是位於檔案中,但在加載期間未被使用的名稱。

將給定的檔案名稱加載到 torch 模型中。此方法專門用於避免 safetensors 不允許的張量共享問題。關於張量共享的更多資訊

safetensors.torch.save_model

< >

( model: Module filename: str metadata: typing.Optional[typing.Dict[str, str]] = None force_contiguous: bool = True )

參數

  • model (torch.nn.Module) — 要保存到磁碟的模型。
  • filename (str) — 要保存檔案的檔案名稱位置。
  • metadata (Dict[str, str], 選填) — 與檔案一起保存的額外資訊。會為每個被丟棄的張量添加一些元數據。這些資訊不足以恢復整個共享結構,但可能有助於理解相關內容。
  • force_contiguous (boolean, 選填, 預設為 True) — 強制將 state_dict 保存為連續的張量。這對模型的正確性沒有影響,但如果當初選擇張量佈局是出於特定原因,這可能會改變效能。

將給定的 torch 模型保存到指定的檔案名稱。此方法專門用於避免 safetensors 不允許的張量共享問題。關於張量共享的更多資訊

在 GitHub 上更新

© . This site is unofficial and not affiliated with Hugging Face, Inc.