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)