Configuration for entropy constraint stage.
| 120 | |
| 121 | @dataclass |
| 122 | class EntropyConfig: |
| 123 | """Configuration for entropy constraint stage.""" |
| 124 | |
| 125 | enabled: bool = False |
| 126 | model_type: Literal["factorized_model", "gaussian_model"] = "factorized_model" |
| 127 | steps: Dict[str, int] = field(default_factory=_default_entropy_steps) |
| 128 | factorized_lr: float = 1e-4 |
| 129 | gaussian_lr: float = 5e-3 |
| 130 | scheduler_gamma: float = 0.01 |
| 131 | rd_lambda: float = 0.01 # Rate-distortion trade-off parameter |
| 132 | |
| 133 | def ensure_all_attributes(self) -> None: |
| 134 | for name in ATTRIBUTE_NAMES: |
| 135 | if name not in self.steps: |
| 136 | self.steps[name] = -1 |
| 137 | |
| 138 | |
| 139 | @dataclass |
no outgoing calls