Class ArgMaxFusion
Defined in File arg_max_fusion.h
Inheritance Relationships
Base Type
public mindspore::ops::BaseOperator
Class Documentation
-
class ArgMaxFusion : public mindspore::ops::BaseOperator
ArgMaxFusion defined ArgMax operator prototype of lite.
Public Functions
-
inline ArgMaxFusion()
Constructor.
-
void Init(const bool keep_dims, const bool out_max_value, const int64_t top_k, const int64_t axis = -1)
Method to init the op's attributes.
- 参数
keep_dims – [in] Define a boolean value to indicate the dimension of output is equal to that of input or not.
out_max_value – [in] Define a boolean value to indicate whether to output maximum value.
top_k – [in] Define the number of maximum value along with axis.
axis – [in] Define where the argmax operation applies to.
-
void set_keep_dims(const bool keep_dims)
Method to set keep_dims attribute.
- 参数
keep_dims – [in] Define a boolean value to indicate the dimension of output is equal to that of input or not.
-
void set_out_max_value(const bool out_max_value)
Method to set out_max_value attribute.
- 参数
out_max_value – [in] Define a boolean value to indicate whether to output maximum value.
-
void set_top_k(const int64_t top_k)
Method to set top_k attribute, default is 1.
- 参数
top_k – [in] Define the number of maximum value along with axis.
-
bool get_keep_dims() const
Method to get keep_dims attribute.
- 返回
a boolean value to indicate the dimension of output is equal to that of input or not.
-
bool get_out_max_value() const
Method to get out_max_value attribute.
- 返回
a boolean value to indicate whether to output maximum value.
-
int64_t get_top_k() const
Method to get top_k attribute.
- 返回
the number of maximum value along with axis.
-
void set_axis(const int64_t axis)
Set axis.
-
void set_output_type(const TypeId output_type)
Set output_type.
-
int64_t get_axis() const
Get axis.
- 返回
axis.
-
TypeId get_output_type() const
Get output_type.
- 返回
output_type.
-
inline ArgMaxFusion()