PQuantML/src/pquant/data_models/pruning_model.py at 3d06f40da3f6a6b70ddc14977cc3262823a00a6b · cern-nextgen/PQuantML · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
from enum import Enum
from typing import List, Literal, Optional

from pydantic import BaseModel, Field


class BasePruningModel(BaseModel):
    disable_pruning_for_layers: List[str] = Field(default_factory=list)
    enable_pruning: bool = Field(default=True)
    threshold_decay: float = Field(default=0.0)


class CSPruningModel(BasePruningModel):
    pruning_method: Literal["cs"] = "cs"
    final_temp: int = Field(default=200)
    threshold_init: float = Field(default=0)


class DSTPruningModel(BasePruningModel):
    pruning_method: Literal["dst"] = "dst"
    alpha: float = Field(default=5.0e-06)
    max_pruning_pct: float = Field(default=0.99)
    threshold_init: float = Field(default=0.0)
    threshold_type: str = Field(default="channelwise")


class FITCompressPruningModel(BasePruningModel):
    pruning_method: Literal["fitcompress"] = "fitcompress"
    min_frac_bits: float = Field(default=2.0)


class PDPPruningModel(BasePruningModel):
    pruning_method: Literal["pdp"] = "pdp"
    epsilon: float = Field(default=0.015)
    sparsity: float = Field(default=0.8)
    temperature: float = Field(default=1.0e-05)
    structured_pruning: bool = Field(default=False)


class WandaPruningModel(BasePruningModel):
    pruning_method: Literal["wanda"] = "wanda"
    M: Optional[int] = (Field(default=None),)
    N: Optional[int] = (Field(default=None),)
    sparsity: float = Field(default=0.9)
    t_delta: int = Field(default=100)
    t_start_collecting_batch: int = Field(default=100)
    calculate_pruning_budget: bool = Field(default=True)


class AutoSparsePruningModel(BasePruningModel):
    pruning_method: Literal["autosparse"] = "autosparse"
    alpha: float = Field(default=0.5)
    alpha_reset_epoch: int = Field(default=90)
    autotune_epochs: int = Field(default=10)
    backward_sparsity: bool = Field(default=False)
    threshold_init: float = Field(default=-5.0)
    threshold_type: str = Field(default="channelwise")


class ActivationPruningModel(BasePruningModel):
    pruning_method: Literal["activation_pruning"] = "activation_pruning"
    threshold: float = Field(default=0.3)
    t_delta: int = Field(default=50)
    t_start_collecting_batch: int = Field(default=50)


class MetricType(str, Enum):
    UNSTRUCTURED = "UnstructuredSparsity"
    STRUCTURED = "StructuredSparsity"
    FPGA_AWARE = "FPGAAwareSparsity"
    PACA_PATTERN = "PACAPatternSparsity"


class ConstraintType(str, Enum):
    EQUALITY = "Equality"
    LEQ = "LessThanOrEqual"
    GEQ = "GreaterThanOrEqual"


class MDMMPruningModel(BasePruningModel):
    pruning_method: Literal["mdmm"] = "mdmm"
    constraint_type: ConstraintType = Field("Equality")
    target_value: float = Field(default=0.0)
    metric_type: MetricType = Field(default="UnstructuredSparsity")
    target_sparsity: float = Field(default=0.9)
    rf: int = Field(default=1)
    epsilon: float = Field(default=1.0e-03)
    scale: float = Field(default=10.0)
    damping: float = Field(default=1.0)
    use_grad: bool = Field(default=False)
    l0_mode: Literal["coarse", "smooth"] = Field(default="coarse")
    scale_mode: Literal["mean", "sum"] = Field(default="mean")
    constraint_lr: float = Field(default=1.0e-3)
    # FPGAAwareSparsityMetric (only used when metric_type == "FPGAAwareSparsity")
    precision: Optional[int] = Field(default=None)
    target_resource: Optional[Literal["DSP", "BRAM"]] = Field(default=None)
    bram_width: Optional[int] = Field(default=None)
    # PACAPatternMetric (only used when metric_type == "PACAPatternSparsity")
    num_patterns_to_keep: Optional[int] = Field(default=None)
    beta: Optional[float] = Field(default=None)
    distance_metric: Optional[Literal["hamming", "valued_hamming", "cosine"]] = Field(default=None)