Safetensors 文件

Torch API

您正在檢視的是需要從原始碼安裝。如果你想使用常規的 pip install,請檢視最新的穩定版本 (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 (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.