Illegal "IsScalar" node encountered in exported ONNX model

[yuanyao-nv]

An exported model contains an IsScalar op which is not in the ONNX spec.

Seems like it comes from the use of common_ops.IsScalar in aten_mean_dim.

Env:
torch 2.3.0.dev20240202+cu121
onnxscript 0.1.0.dev20240203
MONAI installed with pip3 install git+https://github.com/Project-MONAI/MONAI@0.8.1#egg=MONAI

[justinchuby]

IsScalar is a function which should be part of the model local functions. How was the model exported?

[yuanyao-nv]

This is the export script:

import torch
from monai.networks.nets import DenseNet

def DenseNetDataloader():
    yield torch.randn(1,1,128,128).to("cuda")

model = lambda : DenseNet(
            init_features=64,
            growth_rate=32,
            block_config=[6, 12, 24, 16],
            spatial_dims=2,
            in_channels=1,
            out_channels=15
        )
data = next(DenseNetDataloader())

model = model().eval().to('cuda')

export_output = torch.onnx.dynamo_export(
    model,
    data,
)
export_output.save('Clara_DenseNet_dynamo.onnx')

[justinchuby]

Was the error raised from the onnx checker?

[justinchuby]

Could you attach the exported ONNX file?

[yuanyao-nv]

ONNX checker actually doesn't give any error. The error was raised by TensorRT.

I'm checking with the relevant parties whether I can directly share the model. Did the export script I attached above work for you? I used the NGC pytorch container docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:xx.xx-py3 and installed torch, onnxscript and MONAI as I noted in the description. You might also need to install some dependencies manually:

pip install hydra-core sentencepiece youtokentome transformers inflect pyannote.core pyannote.metrics editdistance jiwer h5py sacremoses pytorch-lightning webdataset ijson wandb

[justinchuby]

One thing I would check is if IsScalar is defined as one of the model local functions.

[yuanyao-nv]

See the model attached here: https://drive.google.com/drive/folders/1f3O9Z2zOTFGlh1grY1l2vQWxE7koQb2j?usp=drive_link
It seems like IsScalar is a node rather than a local function. It's location is
torch_nn_modules_container_Sequential_class_layers_1, torch_nn_modules_pooling_AdaptiveAvgPool2d_class_layers_pool_1, aten_mean_dim.

[justinchuby]

Looks like the gdrive folder is empty

[justinchuby]

IsScalar should be defined:

import onnx
m = onnx.load("Clara_DenseNet_dynamo.onnx")
print(m.functions[-1])

yields

name: "IsScalar"
input: "input"
output: "return_val"
node {
  input: "input"
  output: "tmp"
  name: "n0"
  op_type: "Shape"
  domain: ""
}
node {
  input: "tmp"
  output: "tmp_0"
  name: "n1"
  op_type: "Size"
  domain: ""
}
node {
  output: "tmp_1"
  name: "n2"
  op_type: "Constant"
  attribute {
    name: "value_int"
    i: 0
    type: INT
  }
  domain: ""
}
node {
  input: "tmp_0"
  input: "tmp_1"
  output: "return_val"
  name: "n3"
  op_type: "Equal"
  domain: ""
}
doc_string: "Return whether the input has rank 0, or is a scalar."
opset_import {
  domain: ""
  version: 18
}
domain: "pkg.onnxscript.torch_lib.common"

But netron does not show the function definition of it. 🤔 I wonder what went wrong?

[yuanyao-nv]

Also I noticed if run the inliner and then onnx-rewriter on this model, the optimized graph still contains the IsScalar node:

[justinchuby]

That's an issue. @gramalingam @xadupre @BowenBao do you have any guesses?

[BowenBao]

I have also encountered the inliner issue around IsScalar. The node is not inlined yet its FunctionProto is removed.