Safetensors
您正在檢視的是需要從原始碼安裝。如果你想使用常規的 pip install,請檢視最新的穩定版本 (
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
, or os.PathLike
) — 包含張量的檔名 - device (
Union[str, int]
, optional, defaults to cpu
) — 載入後張量需要位於的裝置。可用選項是所有常規的 torch 裝置位置。
返回
Dict[str, torch.Tensor]
一個字典,其中鍵為名稱,值為 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]
一個字典,其中鍵為名稱,值為 CPU 上的 torch.Tensor
從純位元組資料載入 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]
) — 輸入的張量。張量需要是連續且密集的。 - filename (
str
, or os.PathLike
)) — 我們要儲存到的檔名。 - metadata (
Dict[str, str]
, optional, defaults to None
) — 您可能想要儲存在檔案頭中的可選純文字元資料。例如,它可以用於指定更多關於底層張量的資訊。這純粹是資訊性的,不影響張量載入。
將張量字典以 safetensors 格式儲存為原始位元組。
示例
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]
) — 輸入的張量。張量需要是連續且密集的。 - metadata (
Dict[str, str]
, optional, defaults to 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
, or os.PathLike
) — 要從中載入檔案的檔名位置。 - strict (
bool
, optional, defaults to True) — 當存在缺失鍵或意外部索引鍵時是否失敗。當為 false 時,函式僅返回缺失和意外的名稱。 - device (
Union[str, int]
, optional, defaults to 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]
, optional) — 與檔案一起儲存的額外資訊。將為每個被丟棄的張量新增一些元資料。這些資訊不足以恢復整個共享結構,但可能有助於理解。 - force_contiguous (
boolean
, optional, defaults to True) — 強制將 state_dict 儲存為連續張量。這對模型的正確性沒有影響,但如果張量的佈局是為此特定原因選擇的,則可能會改變效能。
將給定的 torch 模型儲存到指定的檔名。此方法專門用於避免 safetensors
中不允許的張量共享問題。關於張量共享的更多資訊
< > 在 GitHub 上更新
© . This site is unofficial and not affiliated with Hugging Face, Inc.