如果无需对模型中的所有子模块进行量化,则可使用模块部分量化。除使用通用 vai_q_pytorch API 外,您也可使用成对 QuantStub/DeQuantStub
运算符来调整其大小。以下示例演示了如何量化 subm0
和 subm2
,而不量化 subm1
。
from pytorch_nndct.nn import QuantStub, DeQuantStub
class WholeModule(torch.nn.module):
def __init__(self,...):
self.subm0 = ...
self.subm1 = ...
self.subm2 = ...
# define QuantStub/DeQuantStub submodules
self.quant = QuantStub()
self.dequant = DeQuantStub()
def forward(self, input):
input = self.quant(input) # begin of part to be quantized
output0 = self.subm0(input)
output0 = self.dequant(output0) # end of part to be quantized
output1 = self.subm1(output0)
output1 = self.quant(output1) # begin of part to be quantized
output2 = self.subm2(output1)
output2 = self.dequant(output2) # end of part to be quantized