()
| 35 | |
| 36 | |
| 37 | def test_qat_convbn2d(): |
| 38 | in_channels = 32 |
| 39 | out_channels = 64 |
| 40 | kernel_size = 3 |
| 41 | |
| 42 | class TestNet(Module): |
| 43 | def __init__(self, groups, bias): |
| 44 | super().__init__() |
| 45 | self.quant = QuantStub() |
| 46 | self.dequant = DequantStub() |
| 47 | self.conv_bn = ConvBn2d( |
| 48 | in_channels, out_channels, kernel_size, groups=groups, bias=bias, |
| 49 | ) |
| 50 | |
| 51 | def forward(self, inp): |
| 52 | out = self.quant(inp) |
| 53 | out = self.conv_bn(out) |
| 54 | out = self.dequant(out) |
| 55 | return out |
| 56 | |
| 57 | inputs = tensor(np.random.randn(4, in_channels, 32, 32).astype(np.float32)) |
| 58 | for groups, bias in product([1, 4], [True, False]): |
| 59 | net = TestNet(groups, bias) |
| 60 | net.train() |
| 61 | qat_net = quantize_qat(net, inplace=False) |
| 62 | disable_fake_quant(qat_net) |
| 63 | normal_outputs = net(inputs) |
| 64 | qat_outputs = qat_net(inputs) |
| 65 | np.testing.assert_allclose( |
| 66 | normal_outputs.numpy(), qat_outputs.numpy(), atol=1e-4, |
| 67 | ) |
| 68 | np.testing.assert_allclose( |
| 69 | net.conv_bn.bn.running_mean.numpy(), |
| 70 | qat_net.conv_bn.bn.running_mean.numpy(), |
| 71 | atol=5e-8, |
| 72 | ) |
| 73 | np.testing.assert_allclose( |
| 74 | net.conv_bn.bn.running_var.numpy(), |
| 75 | qat_net.conv_bn.bn.running_var.numpy(), |
| 76 | atol=5e-7, |
| 77 | ) |
| 78 | net.eval() |
| 79 | normal_outputs = net(inputs) |
| 80 | qat_net.eval() |
| 81 | qat_outputs = qat_net(inputs) |
| 82 | np.testing.assert_allclose( |
| 83 | normal_outputs.numpy(), qat_outputs.numpy(), atol=1e-4, |
| 84 | ) |
| 85 | |
| 86 | |
| 87 | def test_qat_convtransposebn2d(): |
nothing calls this directly
no test coverage detected