Safetensors
您目前正在檢視
main 版本,這需要從原始碼安裝. 如果您希望使用一般的 pip 安裝,請查看最新的穩定版本 (
v0.5.0-rc.0)。
加入 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 (
str 或 os.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 (
str 或 os.PathLike) — 我們要保存到的檔案名稱。 - metadata (
Dict[str, str], 選填, 預設為 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) — 您可能想在標頭中保存的僅文字元數據。例如,這對於指定關於底層張量的更多資訊很有用。這純粹是資訊性的,不會影響張量的加載。
代表該格式的原始位元組。
將張量字典保存為 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 (
str 或 os.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.