sparsify

torchao.sparsity.sparsify_(model: Module, apply_tensor_subclass: Callable[[Tensor], Tensor], filter_fn: Optional[Callable[[Module, str], bool]] = None) → Module

使用 apply_tensor_subclass 轉換模型中線性模組的權重。此函式本質上與 quantize 相同，但用於稀疏性子類。

目前，我們支援三種稀疏性選項

• 使用 semi_sparse_weight 的半結構化 (2:4) 稀疏性

• 使用 layout=SemiSparseLayout 的 int8 動態量化 + 2:4 稀疏性

• 使用 layout=SparseMarlinLayout 的 int4 僅權重量化 + 2:4 稀疏性

引數：

• model (torch.nn.Module) – 輸入模型

• apply_tensor_subclass (Callable[[torch.Tensor], torch.Tensor]) – 將浮點 Tensor 轉換為（稀疏化後的）張量子類例項的函式（例如仿射量化張量例項）

• filter_fn (Optional[Callable[[torch.nn.Module, str], bool]]) – 一個函式，接受 nn.Module 例項和模組的完全限定名稱，如果想要在模組的權重上執行 apply_tensor_subclass，則返回 True

示例

import torch
import torch.nn as nn
from torchao.sparsity import sparsify_

def filter_fn(module: nn.Module, fqn: str) -> bool:
    return isinstance(module, nn.Linear)

m = nn.Sequential(nn.Linear(32, 1024), nn.Linear(1024, 32))

# for 2:4 sparsity
from torchao.sparse_api import semi_sparse_weight
m = sparsify_(m, semi_sparse_weight(), filter_fn)

# for int8 dynamic quantization + 2:4 sparsity
from torchao.dtypes import SemiSparseLayout
m = quantize_(m, int8_dynamic_activation_int8_weight(layout=SemiSparseLayout), filter_fn)