| 422 | |
| 423 | |
| 424 | class MovingAverageAbsMaxScale(Layer): |
| 425 | def __init__( |
| 426 | self, |
| 427 | name: str | None = None, |
| 428 | moving_rate: float = 0.9, |
| 429 | dtype: DTypeLike = 'float32', |
| 430 | reduce_type: Literal['max'] | None = None, |
| 431 | ) -> None: |
| 432 | r""" |
| 433 | MovingAverageMaxScale layer is used to calculating the output quantization |
| 434 | scale of Layer. Its computational formula is described as below: |
| 435 | |
| 436 | :math:`scale = (moving\_rate*accum+max(abs(x)))/(moving\_rate*state+1)` |
| 437 | :math:`Out = X` |
| 438 | """ |
| 439 | super().__init__() |
| 440 | self._moving_rate = moving_rate |
| 441 | self._reduce_type = reduce_type |
| 442 | scale_prefix = f'{name}.scale' if name else 'outscale.scale' |
| 443 | scale_name = unique_name.generate(scale_prefix) |
| 444 | scale_attr = ParamAttr( |
| 445 | name=scale_name, initializer=Constant(0), trainable=False |
| 446 | ) |
| 447 | self._scale = self.create_parameter( |
| 448 | shape=[1], attr=scale_attr, dtype=dtype |
| 449 | ) |
| 450 | self._scale.stop_gradient = True |
| 451 | |
| 452 | state_prefix = f"{name}.state" if name else 'outscale.state' |
| 453 | state_attr = ParamAttr( |
| 454 | name=unique_name.generate(state_prefix), |
| 455 | initializer=Constant(0), |
| 456 | trainable=False, |
| 457 | ) |
| 458 | self._state = self.create_parameter( |
| 459 | shape=[1], attr=state_attr, dtype=dtype |
| 460 | ) |
| 461 | self._state.stop_gradient = True |
| 462 | |
| 463 | accum_prefix = f"{name}.accum" if name else 'outscale.accum' |
| 464 | accum_attr = ParamAttr( |
| 465 | name=unique_name.generate(accum_prefix), |
| 466 | initializer=Constant(0), |
| 467 | trainable=False, |
| 468 | ) |
| 469 | self._accum = self.create_parameter( |
| 470 | shape=[1], attr=accum_attr, dtype=dtype |
| 471 | ) |
| 472 | self._accum.stop_gradient = True |
| 473 | |
| 474 | def forward(self, input: Tensor) -> Tensor: |
| 475 | if in_dynamic_mode(): |
| 476 | attrs = ( |
| 477 | 'moving_rate', |
| 478 | self._moving_rate, |
| 479 | 'is_test', |
| 480 | not self.training, |
| 481 | ) |