()
| 337 | |
| 338 | |
| 339 | def test_qat_linearbn1d(): |
| 340 | in_features = 10 |
| 341 | out_features = 5 |
| 342 | |
| 343 | class TestNet(Module): |
| 344 | def __init__(self, bias): |
| 345 | super().__init__() |
| 346 | self.quant = QuantStub() |
| 347 | self.dequant = DequantStub() |
| 348 | self.linear_bn = LinearBn1d(in_features, out_features, bias=bias,) |
| 349 | |
| 350 | def forward(self, inp): |
| 351 | out = self.quant(inp) |
| 352 | out = self.linear_bn(out) |
| 353 | out = self.dequant(out) |
| 354 | return out |
| 355 | |
| 356 | inputs = tensor(np.random.randn(4, in_features).astype(np.float32)) |
| 357 | for bias in [True, False]: |
| 358 | net = TestNet(bias) |
| 359 | net.train() |
| 360 | qat_net = quantize_qat(net, inplace=False) |
| 361 | disable_fake_quant(qat_net) |
| 362 | normal_outputs = net(inputs) |
| 363 | qat_outputs = qat_net(inputs) |
| 364 | np.testing.assert_allclose( |
| 365 | normal_outputs.numpy(), qat_outputs.numpy(), atol=1e-6, |
| 366 | ) |
| 367 | np.testing.assert_allclose( |
| 368 | net.linear_bn.bn.running_mean.numpy(), |
| 369 | qat_net.linear_bn.bn.running_mean.numpy(), |
| 370 | atol=5e-8, |
| 371 | ) |
| 372 | np.testing.assert_allclose( |
| 373 | net.linear_bn.bn.running_var.numpy(), |
| 374 | qat_net.linear_bn.bn.running_var.numpy(), |
| 375 | atol=5e-7, |
| 376 | ) |
| 377 | |
| 378 | net.eval() |
| 379 | normal_outputs = net(inputs) |
| 380 | qat_net.eval() |
| 381 | qat_outputs = qat_net(inputs) |
| 382 | np.testing.assert_allclose( |
| 383 | normal_outputs.numpy(), qat_outputs.numpy(), atol=1e-6, |
| 384 | ) |
| 385 | |
| 386 | |
| 387 | def test_qat_linear_relu(): |
nothing calls this directly
no test coverage detected