[Bug] fix setitem with scalar or single element tensor

mmdeploy/pytorch/functions/tensor_setitem.py

@@ -4,7 +4,7 @@
 import torch
 from packaging.version import parse
 
-from mmdeploy.core import FUNCTION_REWRITER
+from mmdeploy.core import FUNCTION_REWRITER, SYMBOLIC_REWRITER
 
 
 @FUNCTION_REWRITER.register_rewriter(func_name='torch.Tensor.__setitem__')
@@ -26,6 +26,22 @@ def tensor__setitem__default(ctx, self, key, value):
  return ctx.origin_func(self, key, value)
 
  out = value
+
+ # value could be scalar or single value Tensor
+ self_shape = self.shape
+ out_shape = [0] * len(self_shape)
+ for i, k in enumerate(key):
+ start = 0 if k.start is None else k.start
+ start = start if start >= 0 else self_shape[i] + start
+ stop = self_shape[i] if k.stop is None else k.stop
+ stop = stop if stop >= 0 else self_shape[i] + stop
+ out_shape[i] = stop - start
+
+ if not isinstance(out, torch.Tensor):
+ out = self.new_full(out_shape, out)
+ elif out.numel() == 1:
+ out = out.expand(out_shape)
+
  for i, k in enumerate(key):
  if k == slice(None):
  continue
@@ -55,3 +71,10 @@ def tensor__setitem__default(ctx, self, key, value):
  # self assign
  # Note that set item does not return any value
  self[...] = out
+
+
+if parse(torch.__version__) >= parse('1.12.0'):
+
+ @SYMBOLIC_REWRITER.register_symbolic('copy', is_pytorch=True)
+ def copy__default(ctx, g, x, y, non_blocking):
+ return x

tests/test_pytorch/test_pytorch_functions.py

@@ -343,6 +343,37 @@ def setitem_slice(x, y):
  assert node.op_type != 'ScatterND'
 
 
+@backend_checker(Backend.ONNXRUNTIME)
+@pytest.mark.skipif(
+ parse(torch.__version__) < parse('1.9.0'), reason='requires torch>1.8.0')
+@pytest.mark.parametrize('x', [torch.rand(1, 3, 16, 16)])
+def test_tensor_setitem_scalar(x):
+ import onnx
+
+ from mmdeploy.utils.test import get_onnx_model
+
+ def setitem_slice(x):
+ H, W = x.shape[-2:]
+ x[:, :, 2:H - 2, 2:W - 2] = 1
+ x[:, :, 4:H - 4, 4:W - 4] = x.new_tensor(2)
+ return x
+
+ wrapped_func = WrapFunction(setitem_slice)
+ model_inputs = {'x': x}
+
+ deploy_cfg = mmcv.Config(
+ dict(
+ onnx_config=dict(input_shape=None),
+ backend_config=dict(type='onnxruntime'),
+ codebase_config=dict(type='mmdet', task='ObjectDetection')))
+ ir_file_path = get_onnx_model(wrapped_func, model_inputs, deploy_cfg)
+
+ onnx_model = onnx.load(ir_file_path)
+ nodes = onnx_model.graph.node
+ for node in nodes:
+ assert node.op_type != 'ScatterND'
+
+
 @pytest.mark.parametrize('output_size', [1, 3])
 def test_adaptive_avg_pool2d(output_size):
  input = torch.rand(1, 3, 6, 6)