sparsify

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

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

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

• 使用 semi_sparse_weight 進行半結構化（2:4）稀疏

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

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

引數:

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

• config (AOBaseConfig) – 工作流配置物件

• filter_fn (Optional[Callable[[torch.nn.Module, str], bool]]) – 一個函式，它接受一個 nn.Module 例項和該模組的完全限定名稱，如果我們要將指定的工作流應用於此模組，則返回 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)