Is the Gradient Operator already supported?

[tiberiusferreira]

Hello, is https://github.com/onnx/onnx/blob/master/docs/Operators.md#aionnxpreviewtraininggradient already supported?

I'm trying to create a toy front-end that allows creation of training graphs and then uses ONNX for the training itself: https://github.com/tiberiusferreira/onnx-experiment

But I've ran into a problem trying to use the Gradient operator. I want to use this operator so I get more control over the training loop.

I have compiled ONNX runtime from source with training enabled and get the error:

onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : Fatal error: Gradient is not a registered function/op

when running this graph:

ir_version: 7
producer_name: "rust"
producer_version: "0.1.0"
graph {
  node {
    input: "left"
    input: "right"
    output: "out"
    op_type: "Add"
  }
  node {
    input: "left"
    output: "deriv"
    op_type: "Gradient"
    attribute {
      name: "xs"
      strings: "left"
      type: STRINGS
    }
    attribute {
      name: "y"
      strings: "out"
      type: STRINGS
    }
    domain: "ai.onnx.preview.training"
  }
  name: "rust-graph"
  initializer {
    dims: 2
    dims: 2
    data_type: 1
    float_data: 1
    float_data: 2
    float_data: 3
    float_data: 4
    name: "left"
  }
  initializer {
    dims: 2
    dims: 2
    data_type: 1
    float_data: 1
    float_data: 2
    float_data: 3
    float_data: 4
    name: "right"
  }
  input {
    name: "left"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 2
          }
          dim {
            dim_value: 2
          }
        }
      }
    }
  }
  input {
    name: "right"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 2
          }
          dim {
            dim_value: 2
          }
        }
      }
    }
  }
  output {
    name: "out"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 2
          }
          dim {
            dim_value: 2
          }
        }
      }
    }
  }
  output {
    name: "deriv"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 2
          }
          dim {
            dim_value: 2
          }
        }
      }
    }
  }
  value_info {
    name: "out"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 2
          }
          dim {
            dim_value: 2
          }
        }
      }
    }
  }
  value_info {
    name: "deriv"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 2
          }
          dim {
            dim_value: 2
          }
        }
      }
    }
  }
}
opset_import {
  version: 12
}
opset_import {
  domain: "ai.onnx.preview.training"
  version: 1
}

with this code:

model = onnx.load_model("my_model.onnx")

model_desc = {
    "inputs": [ 
        ("left", ["batch", "1"])
    ],
    "outputs": [('deriv', [], True)],
}
options = {'device': {'id': 'CPU'}}
trainer = ORTTrainer(model, model_desc, optim_config=SGDConfig(), options=ORTTrainerOptions(options))

Also, is there a way to use onnxruntime.training.TrainingSession to just run a graph with Gradient ops and get the result instead of having to setup a ORTTrainer?

[georgen117]

I am just a user so take my comment with almost no authority.

The gradient ops are specified per-op. There is not a global Gradient operator that I am aware of. I am not great at reading the JSON version of the onnx model but from the looks of your model, the only operator you have is the Add operator. The name of the corresponding gradient operator would be AddGrad.

I don't think you actually need to add the AddGrad or any gradient operator to the model. You only need to provide the inference version of the model and the ORTTrainer will do the magic of generating the per-operator gradients. This is based on looking at the orrttraining/pytorch_frontend_exampls/mnist_training.py. The model specified there are only the forward (inference) operators. This observation is also based on playing with the onnxruntime_training_mnist.cc sample.