torch.index_put conversion error

[RocketFlash]

Hi! I am trying to convert NN with the following operation inside:

accum = torch.zeros((N, C, self.numangle, self.numrho), dtype=torch.float32)

...

accum.index_put_(indices=[n_i_t, c_i_t, angls, accum_idxs],
                                 values=img[n_i_t, c_i_t, y_idxs, x_idxs],
                                 accumulate=True)

jit.trace goes well, but on ct.convert step I've got the following error:

File "/Users/user/opt/anaconda3/envs/retech_dht/lib/python3.8/site-packages/coremltools/converters/mil/mil/ops/defs/iOS15/tensor_operation.py", line 1282, in value_inference
    return np.stack(values, self.axis.val)
  File "<__array_function__ internals>", line 180, in stack
  File "/Users/user/opt/anaconda3/envs/retech_dht/lib/python3.8/site-packages/numpy/core/shape_base.py", line 429, in stack
axis = normalize_axis_index(axis, result_ndim)
numpy.AxisError: axis 3 is out of bounds for array of dimension 2

Conversion performed well if I comment line with indexput

python: 3.8 torch version: 1.12.1 coremltools version: 6.1

[TobyRoseman]

Looking at the code I'm not seeing anything obviously work.

@RocketFlash - There are many undefined variables in your code snippet. Can you give us a standalone example to reproduce this problem?

As a temporary workaround, you could try just deleting the value_inference method.

[RocketFlash]

@TobyRoseman thank you for reply. Here is a standalone example:

import torch
import torch.nn as nn
import coremltools as ct

class Test(nn.Module):
    def __init__(self):
        super(Test, self).__init__()

    def forward(self, img):
        accum = torch.zeros((1, 1, 100, 100), dtype=torch.float32)

        a = torch.zeros((100)).long()
        b = torch.zeros((100)).long()
        c = torch.arange(100).long()
        d = torch.arange(100).long()

        e = torch.ones((100)).long()
        f = torch.ones((100)).long()

        accum.index_put_(indices=[a, b, c, d],
                        values=img[a, b, e, f],
                        accumulate=True)

        return accum

if __name__ == '__main__':
    sample = torch.ones((1, 1, 100, 100), dtype=torch.float32)
    test_model = Test()
    test_model.eval()

    print('Model tracing')
    traced_model = torch.jit.trace(test_model, sample)
    print('Tracing completed')

    print('Convert to coreml')
    cm_model = ct.convert(traced_model,
                                inputs=[ct.TensorType(name="input",
                                                      shape=sample.shape)],
                                outputs=[ct.TensorType(name="output")]
                                )
    print('Convertion completed')

[RocketFlash]

@TobyRoseman sorry for bothering again, I reorginized code, now it works and model converts to .mlmodel, but it doesn't work on inference. Here is a new snippet and model graph:

import torch
import torch.nn as nn
import coremltools as ct
import numpy as np

@torch.jit.script_if_tracing
def accumulator_iteration(img:torch.Tensor, accum:torch.Tensor) -> torch.Tensor:
    idxs = torch.nonzero(img).reshape(-1, 4).long()
    n_idxs = idxs[:, 0]
    c_idxs = idxs[:, 1]
    y_idxs = idxs[:, 2]
    x_idxs = idxs[:, 3]
    n_feat = idxs.size(dim=0)

    if n_feat>0:
        c = torch.ones((n_feat,)).long()
        d = torch.ones((n_feat,)).long()
        accum = torch.index_put(accum, indices=[n_idxs, c_idxs, c, d],
                                       values=img[n_idxs, c_idxs, y_idxs, x_idxs],
                                       accumulate=True)
    return accum

class Test(nn.Module):
    def __init__(self):
        super(Test, self).__init__()

    def forward(self, img):
        accum = torch.zeros((1, 10, 100, 100)).float()
        return accumulator_iteration(img, accum)

if __name__ == '__main__':
    torch.manual_seed(28)
    sample = torch.rand(1, 10, 100, 100).float()
    test_model = Test()
    test_model.eval()

    print('Model tracing')
    traced_model = torch.jit.trace(test_model, sample)
    traced_model.eval()
    print('Tracing completed')

    print(traced_model.code)    
    traced_model.save(f'./weights/test.pt')

    sample_test = torch.rand(1, 10, 100, 100).float()
    o_traced = traced_model(sample_test)
    print(f'traced : {o_traced[0,0,:7,:7]}')

    print('Convert to coreml')
    cm_model = ct.convert(traced_model,
                          inputs=[ct.TensorType(name="input",
                                                shape=sample.shape)],
                          outputs=[ct.TensorType(name="output")],
                          debug=True)
    print('Convertion completed')
    cm_model.save(f"./weights/test.mlmodel")

    o_coreml = cm_model.predict({'input': sample_test.cpu()})['output']
    print(f'coreml : {o_coreml[0,0,:7,:7]}')

It seems coremltools didn't connect model with output:

here is an error I got

Traceback (most recent call last):
  File "test_coreml_snippet.py", line 59, in <module>
    o_coreml = cm_model.predict({'input': sample_test.cpu()})['output']
  File "/Users/raufyagfarov/opt/anaconda3/envs/retech_dht/lib/python3.8/site-packages/coremltools/models/model.py", line 509, in predict
    return self.__proxy__.predict(data)
RuntimeError: {
    NSLocalizedDescription = "Error computing NN outputs.";
}

[psilvio]

@RocketFlash I face the exact same problem with the conversion of a tensorflow segmentation model.

Some more infos for the converted coreML model:

input {
  name: "input"
  shortDescription: "Input image"
  type {
    multiArrayType {
      shape: 1
      shape: 160
      shape: 224
      shape: 192
      shape: 1
      dataType: FLOAT32
      shapeRange {
        sizeRanges {
          lowerBound: 1
          upperBound: -1
        }
        sizeRanges {
          lowerBound: 160
          upperBound: 160
        }
        sizeRanges {
          lowerBound: 224
          upperBound: 224
        }
        sizeRanges {
          lowerBound: 192
          upperBound: 192
        }
        sizeRanges {
          lowerBound: 1
          upperBound: 1
        }
      }
    }
  }
}
output {
  name: "output"
  shortDescription: "Segmentation"
  type {
    multiArrayType {
      dataType: FLOAT32
    }
  }
}
metadata {
  userDefined {
    key: "com.github.apple.coremltools.source"
    value: "tensorflow==2.10.0"
  }
  userDefined {
    key: "com.github.apple.coremltools.version"
    value: "6.1"
  }
}

Prediction results in following output:

Traceback (most recent call last):
  File "~/CoreML/coreml_implementation_mets.py", line 56, in <module>
    results = mlmodel.predict({"input": input_data_mod})['output']
  File "~/miniconda/envs/tf_ml/lib/python3.9/site-packages/coremltools/models/model.py", line 509, in predict
    return self.__proxy__.predict(data)
RuntimeError: {
    NSLocalizedDescription = "Error computing NN outputs.";
}