Function test_tensorrt_int8_calibration

tests/python/relax/test_codegen_tensorrt.py:272–316

(monkeypatch)

def test_tensorrt_int8_calibration(monkeypatch):
    # INT8 calibration path: the first N runs feed calibration batches, then the INT8 engine is
    # built and run. Validates that the calibrator copies a full batch (batch_size * per-sample
    # elements) without over-reading the input or over-writing the device buffers, which previously
    # crashed for batch > 1.
    @tvm.script.ir_module
    class Conv2dInt8:
        @R.function
        def main(
            data: R.Tensor((2, 8, 16, 16), "float32"), weight: R.Tensor((4, 8, 3, 3), "float32")
        ):
            with R.dataflow():
                out = relax.op.nn.conv2d(data, weight, padding=1)
                R.output(out)
            return out

    data = np.random.randn(2, 8, 16, 16).astype("float32")
    weight = np.random.randn(4, 8, 3, 3).astype("float32")
    ref = build_and_run(Conv2dInt8, [data, weight], "llvm", legalize=True)

    patterns = [("tensorrt.nn.conv2d", is_op("relax.nn.conv2d")(wildcard(), wildcard()))]
    offloaded = tvm.transform.Sequential(
        [
            relax.transform.BindParams("main", {"weight": weight}),
            relax.transform.FuseOpsByPattern(patterns),
            relax.transform.MergeCompositeFunctions(),
            relax.transform.RunCodegen(),
        ]
    )(Conv2dInt8)

    num_calibration_batches = 2
    monkeypatch.setenv("TVM_TENSORRT_USE_INT8", "1")
    monkeypatch.setenv("TENSORRT_NUM_CALI_INT8", str(num_calibration_batches))

    dev = tvm.device("cuda", 0)
    vm = relax.VirtualMachine(tvm.compile(offloaded, "cuda"), dev)
    data_trt = tvm.runtime.tensor(data, dev)
    out = None
    for _ in range(num_calibration_batches + 1):
        out = vm["main"](data_trt).numpy()

    assert np.isfinite(out).all()
    # INT8 is lossy, so use a generous tolerance; the key assertion is that calibration completed
    # without a CUDA error.
    tvm.testing.assert_allclose(out, ref, rtol=0.2, atol=0.1 * float(np.abs(ref).max()))


if __name__ == "__main__":