モデル内のすべてのサブモジュールを量子化する必要がない場合は、モジュールの部分的量子化を使用できます。汎用 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