Function test_conv2d

tests/python/nightly/test_nnapi/test_ops.py:291–321

()

def test_conv2d():
    remote_obj, tracker = remote()

    def create_model() -> tvm.IRModule:
        @tvm.script.ir_module
        class Module:
            @R.function
            def main(
                i0: R.Tensor((1, 3, 224, 224), "float32"),
                i1: R.Tensor((64, 3, 3, 3), "float32"),
                i2: R.Tensor((1, 64, 1, 1), "float32"),
            ):
                with R.dataflow():
                    t0 = R.nn.conv2d(i0, i1, strides=(1, 1), padding=(1, 1))
                    t0 = R.add(i2, t0)
                    R.output(t0)
                return t0

        return Module

    mod = create_model()
    verify(
        remote_obj,
        tracker,
        mod,
        inputs=[
            np.random.random(size=(1, 3, 224, 224)).astype("float32"),
            np.random.random(size=(64, 3, 3, 3)).astype("float32"),
            np.random.random(size=(1, 64, 1, 1)).astype("float32"),
        ],
    )


def test_max_pool2d():