gm.nn.QuantizationAwareWrapper

class gemma.gm.nn.QuantizationAwareWrapper(

*,

method: gemma.peft._quantization_utils.QuantizationMethod = QuantizationMethod.NONE,

model: flax.linen.module.Module,

parent: flax.linen.module.Module | flax.core.scope.Scope | flax.linen.module._Sentinel | None = <flax.linen.module._Sentinel object>,

name: str | None = None,

)

基类：flax.linen.module.Module

围绕 Gemma 模型进行封装，以启用量化感知训练。

封装的模型将使其所有 nn.Dense、nn.Einsum 等层替换为其量化感知训练版本。 有关更多详细信息，请参阅 gemma.peft 文档。

method

要使用的量化方法。

类型:

gemma.peft._quantization_utils.QuantizationMethod

model

要封装的模型。

类型:

flax.linen.module.Module

method: gemma.peft._quantization_utils.QuantizationMethod = 'none'

model: flax.linen.module.Module

name: str | None = None

parent: flax.linen.module.Module | flax.core.scope.Scope | flax.linen.module._Sentinel | None = None

scope: Scope | None = None