From 29f7cdb84a20cb3a6b1d8699340b829cb78bd5b3 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Tue, 8 Jun 2021 15:43:59 +0800 Subject: [PATCH] fix(mgb/opr): correct nvof out shape computation GitOrigin-RevId: 16bf086e92125cba867d9b935bb487363602014e --- imperative/python/megengine/functional/vision.py | 3 ++- src/opr/impl/misc.cpp | 5 +++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/imperative/python/megengine/functional/vision.py b/imperative/python/megengine/functional/vision.py index 5e3bcfb1..f1427cee 100644 --- a/imperative/python/megengine/functional/vision.py +++ b/imperative/python/megengine/functional/vision.py @@ -633,7 +633,8 @@ def nvof(src: Tensor, precision: int = 1) -> Tensor: :src shape: input tensor with shape (n, t, h, w, c4). :src dtype: uint8. :param precision: 0:NV_OF_PERF_LEVEL_SLOW 1:NV_OF_PERF_LEVEL_MEDIUM 2:NV_OF_PERF_LEVEL_FAST. - :output shape: (n, t-1, h//4, w//4, c2). + :output shape: ``(n, t-1, (h+out_grid_size-1)//out_grid_size, (w+out_grid_size-1)//out_grid_size, c2)``. + By default, out_grid_size = 4. :output dtype: int16. .. code-block:: python diff --git a/src/opr/impl/misc.cpp b/src/opr/impl/misc.cpp index 9d0820d8..558ad0a2 100644 --- a/src/opr/impl/misc.cpp +++ b/src/opr/impl/misc.cpp @@ -224,6 +224,7 @@ void NvOf::scn_do_execute() { void NvOf::init_output_static_infer_desc() { using namespace cg::static_infer; auto infer_shape = [](TensorShape& dest, const InpVal& iv) { + auto out_grid_size = NV_OF_OUTPUT_VECTOR_GRID_SIZE_4; auto ishp = iv.val.at(0).shape(); //! nvof input format: nthwc4 mgb_assert(ishp.ndim == 5); @@ -232,8 +233,8 @@ void NvOf::init_output_static_infer_desc() { SmallVector tv; tv.push_back(ishp[0]); tv.push_back(ishp[1] - 1); - tv.push_back(ishp[2] / 4); - tv.push_back(ishp[3] / 4); + tv.push_back((ishp[2] + out_grid_size - 1) / out_grid_size); + tv.push_back((ishp[3] + out_grid_size - 1) / out_grid_size); tv.push_back(ishp[4] / 2); dest = tv; -- GitLab