def init(self, C, M, N, dtype, mode): assert(mode in ('Q', 'D')) self.input = torch.rand(C, M, N) self.dtype = dtype self.op = nnq.Quantize(scale=1.0, zero_point=0, dtype=dtype) self.set_module_name('QuantizePerTensor') if mode == 'D': self.input = self.op(self.input) self.op = nnq.DeQuantize() self.set_module_name('DequantizePerTensor')
def test_quant_dequant_api(self): r = torch.tensor([[1., -1.], [1., -1.]], dtype=torch.float) scale, zero_point, dtype = 1.0, 2, torch.qint8 # testing Quantize API qr = torch.quantize_linear(r, scale, zero_point, dtype) quant_m = nnq.Quantize(scale, zero_point, dtype) qr2 = quant_m(r) self.assertEqual(qr, qr2) # testing Dequantize API rqr = qr.dequantize() dequant_m = nnq.DeQuantize() rqr2 = dequant_m(qr2) self.assertEqual(rqr, rqr2)
def __init__(self, inputs_to_qparams_map): super(ConditionalQuantize, self).__init__() self.quantizers = nn.ModuleDict() for idx, qparams in inputs_to_qparams_map.items(): self.quantizers[str(idx)] = nnq.Quantize(*qparams)