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)