Caffe2 - C++ API
A deep learning, cross platform ML framework
Public Types | Public Member Functions | Friends
caffe2::OpSchema Class Reference

A class to record the schema of an op. More...

#include <operator_schema.h>

Public Types

typedef std::function< vector< TensorShape >const OperatorDef &, const vector< TensorShape > &)> TensorInferenceFunctionType
 

Public Member Functions

 OpSchema (const string &file, const int line)
 
const string & file () const
 Returns the file that the op schema is registered from.
 
int line () const
 Returns the line in file that the op schema is registered from.
 
const char * doc () const
 Returns the docstring of the op schema.
 
bool Verify (const OperatorDef &def) const
 Verifies if an operator definition protobuf matches the pattern specified in the schema.
 
OpSchemaNumInputs (int n)
 A single input.
 
OpSchemaNumInputs (int min, int max)
 Input could be in range [min, max], inclusive.
 
OpSchemaNumInputs (set< int > allowed_input_nums)
 Input could be one of the values specified in allowed_input_nums.
 
OpSchemaNumInputs (std::function< bool(int)> func)
 Input is checked with a specified function.
 
OpSchemaNumOutputs (int n)
 A single output.
 
OpSchemaNumOutputs (int min, int max)
 Output could be in range [min, max], inclusive.
 
OpSchemaNumOutputs (set< int > allowed_output_nums)
 Output could be one of the values specified in allowed_output_nums.
 
OpSchemaNumOutputs (std::function< bool(int)> func)
 Output is checked with a specified function.
 
OpSchemaNumInputsOutputs (std::function< bool(int, int)> func)
 Relationship between inputs and outputs is checked with a specified function.
 
OpSchemaOutputCalculator (std::function< int(int)> calc)
 Set the output calculator to a user-defined function.
 
OpSchemaSameNumberOfOutput ()
 Set the number of outputs to be the same as the number of inputs.
 
OpSchemaAllowInplace (std::function< bool(int, int)> inplace)
 
OpSchemaAllowInplace (set< std::pair< int, int >> inplace)
 
OpSchemaAllowOneToOneInplace ()
 
OpSchemaEnforceInplace (std::function< bool(int, int)> inplace)
 
OpSchemaEnforceInplace (set< std::pair< int, int >> inplace)
 
OpSchemaEnforceOneToOneInplace ()
 
OpSchemaTensorInferenceFunction (TensorInferenceFunctionType function)
 Sets the tensor inference function, which is a std::function object defined in operator_schema.h.
 
OpSchemaIdenticalTypeAndShape ()
 Seets the tensor inference function to produce the same output as the input.
 
OpSchemaIdenticalTypeAndShapeOfInput (int idx)
 
OpSchemaIdenticalTypeAndShapeOfInputDim (int idx, int dim)
 
OpSchemaScalarType (::caffe2::TensorProto_DataType dt)
 
vector< TensorShape > InferTensor (const OperatorDef &def, const vector< TensorShape > input_type_shape) const
 A function to allow one to infer the type and shape from the op schema.
 
OpSchemaSetDoc (const string &doc)
 
OpSchemaArg (const char *name, const char *description)
 
OpSchemaInput (const int n, const char *name, const char *description)
 
OpSchemaOutput (const int n, const char *name, const char *description)
 
OpSchemaFillUsing (std::function< void(OpSchema &)> populator)
 
int CalculateOutput (int num_input) const
 A function to allow one to get the number of outputs based on the number of inputs, if this schema supports it.
 
const std::vector< std::pair< const char *, const char * > > & arg_desc ()
 
const std::vector< std::pair< const char *, const char * > > & input_desc ()
 
const std::vector< std::pair< const char *, const char * > > & output_desc ()
 

Friends

std::ostream & operator<< (std::ostream &out, const OpSchema &schema)
 

Detailed Description

A class to record the schema of an op.

OpSchema records the common interface of an op specified by its name. This is optional for each operator implemented in Caffe2 but is strongly recommended.

To register an OpSchema, one can use the macro OPERATOR_SCHEMA(name) and then append the various functions in the class. For example, for an op that itakes in two inputs, one output, and the first input and output could be in-place, can be written as

OPERATOR_SCHEMA(name)
    .NumInputs(2).NumOutputs(1).AllowInplace({{0, 0}});

Definition at line 36 of file operator_schema.h.


The documentation for this class was generated from the following files: