From 0df9409634e13c023940534e3f763e61cf3d1a89 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 5 Mar 2021 11:27:57 +0900 Subject: [PATCH] Fix for dynamic batch size conv2d nhwc --- python/tvm/relay/op/strategy/cuda.py | 2 ++ python/tvm/topi/cuda/conv2d_nhwc.py | 4 +++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/python/tvm/relay/op/strategy/cuda.py b/python/tvm/relay/op/strategy/cuda.py index 85bbab692574..3b498a99495e 100644 --- a/python/tvm/relay/op/strategy/cuda.py +++ b/python/tvm/relay/op/strategy/cuda.py @@ -354,6 +354,8 @@ def judge_winograd( OH = (H + pt + pb - KH) // stride_h + 1 OW = (W + pl + pr - KW) // stride_w + 1 nH, nW = (OH + tile_size - 1) // tile_size, (OW + tile_size - 1) // tile_size + if not isinstance(N, int): + return False, False, False P = N * nH * nW judge_winograd_tensorcore = ( diff --git a/python/tvm/topi/cuda/conv2d_nhwc.py b/python/tvm/topi/cuda/conv2d_nhwc.py index a08d217696e2..991585587bbf 100644 --- a/python/tvm/topi/cuda/conv2d_nhwc.py +++ b/python/tvm/topi/cuda/conv2d_nhwc.py @@ -129,4 +129,6 @@ def schedule_conv2d_nhwc_direct(cfg, s, Conv): N, OH, OW, CO = get_const_tuple(output.shape) KH, KW, CI, _ = get_const_tuple(kernel.shape) - cfg.add_flop(2 * N * OH * OW * CO * CI * KH * KW) + + if isinstance(N, int): + cfg.add_flop(2 * N * OH * OW * CO * CI * KH * KW)