Caffe2 - C++ API
A deep learning, cross platform ML framework
Public Types | Public Member Functions | Data Fields | Static Public Attributes
caffe2::AbstractUnsortedSegmentOp< T, SIndex, Context, Reducer, SparseFused, InputAccessor > Class Template Reference

Unsorted segment reduction op with optional fused embedding lookup. More...

#include <segment_reduction_op.h>

Inheritance diagram for caffe2::AbstractUnsortedSegmentOp< T, SIndex, Context, Reducer, SparseFused, InputAccessor >:
caffe2::Operator< Context > caffe2::OperatorBase caffe2::Observable< OperatorBase >

Public Types

enum  { INDICES = Reducer::kInputCount, SEGMENT_IDS = Reducer::kInputCount + (SparseFused ? 1 : 0) }
 
- Public Types inherited from caffe2::Observable< OperatorBase >
using Observer = ObserverBase< OperatorBase >
 

Public Member Functions

template<class... Args>
 AbstractUnsortedSegmentOp (Args &&...args)
 
bool RunOnDevice () override
 
template<typename IndexType >
bool DoRunWithType ()
 
template<typename IndexType , int FixedSize>
bool DoRunWithValue ()
 
- Public Member Functions inherited from caffe2::Operator< Context >
 Operator (const OperatorDef &operator_def, Workspace *ws)
 
 Operator (const c10::FunctionSchema &fn_schema, std::vector< c10::IValue > inputs, std::vector< at::Tensor > outputs)
 
const TensorInput (int idx, DeviceType type=Context::GetDeviceType())
 Retrieve a non-owning reference to the input at position 'idx' for this operator. More...
 
Tensor XOutput (int idx, at::IntArrayRef dims, at::TensorOptions options)
 XOutput is a modernized version of Output which returns a Tensor rather than a Tensor* (the raw pointer in the latter case is useless, as Tensor is a pointer type.)
 
- Public Member Functions inherited from caffe2::OperatorBase
 OperatorBase (const OperatorDef &operator_def, Workspace *ws)
 
 OperatorBase (const c10::FunctionSchema &schema, std::vector< c10::IValue > inputs, std::vector< at::Tensor > outputs)
 
bool isLegacyOperator () const
 Return true if the operator was instantiated with OperatorDef New operators should be instantiated with FunctionSchema.
 
const c10::FunctionSchemagetFunctionSchema () const
 
bool HasArgument (const string &name) const
 Checks if the operator has an argument of the given name.
 
template<typename T >
T GetSingleArgument (const string &name, const T &default_value) const
 
template<typename T >
bool HasSingleArgumentOfType (const string &name) const
 
template<typename T >
vector< TGetVectorFromIValueList (const c10::IValue &value) const
 
template<typename T >
vector< TGetRepeatedArgument (const string &name, const vector< T > &default_value={}) const
 
template<typename T >
const TInput (int idx)
 
template<typename T >
const TInput (int idx, DeviceType type)
 
template<typename T >
TOutput (int idx)
 
template<typename T >
TOutput (int idx, DeviceType type)
 
Tensor XOutputTensor (int idx, at::IntArrayRef dims, at::TensorOptions options)
 
void SetOutputTensor (int idx, Tensor tensor)
 
Tensor OutputTensorOrUndefined (int idx)
 
TensorOutputTensor (int idx, at::IntArrayRef dims, at::TensorOptions options)
 
TensorOutputTensorCopyFrom (int idx, at::TensorOptions options, const Tensor &src, bool async=false)
 
TensorOutputTensorAlias (int idx, const Tensor &src)
 
template<typename T >
TOutput (int idx, T *allocated)
 
const BlobInputBlob (int idx)
 
BlobOutputBlob (int idx)
 
bool IsInputOutputAlias (int i, int j)
 
template<typename T >
bool InputIsType (int idx)
 
bool InputIsTensorType (int idx, DeviceType device_type)
 
template<typename T >
bool OutputIsType (int idx)
 
bool OutputIsTensorType (int idx, DeviceType type)
 
int InputSize () const
 
int OutputSize () const
 
const vector< const Blob * > & Inputs () const
 
const vector< Blob * > & Outputs ()
 
vector< TensorShape > InputTensorShapes () const
 
virtual void WaitEvent (const Event &ev, int=-1)
 
void Wait (const OperatorBase &other, int stream_id=-1)
 
virtual void WaitEvents (const std::vector< const Event * > &events, int=-1)
 
virtual void Finish ()
 
virtual bool Run (int=0)
 
virtual bool HasAsyncPart () const
 
virtual bool SupportsAsyncScheduling () const
 
virtual bool RunAsync (int stream_id=0)
 
virtual void AddRelatedBlobInfo (EnforceNotMet *err)
 
const OperatorDef & debug_def () const
 
void set_debug_def (const std::shared_ptr< const OperatorDef > &operator_def)
 
bool has_debug_def () const
 
void RecordLastFailedOpNetPosition ()
 
int net_position () const
 
void set_net_position (int idx)
 
const DeviceOption & device_option () const
 
const Eventevent () const
 
Eventevent ()
 
void ResetEvent ()
 
void DisableEvent ()
 
bool IsEventDisabled () const
 
virtual void SyncDeviceBarrierForObservers ()
 
virtual bool IsStreamFree (int) const
 
const std::string & type () const
 
void annotate_engine (const std::string &engine)
 
const std::string & engine () const
 
void SetExecutorHelper (ExecutorHelper *helper)
 
ExecutorHelperGetExecutorHelper () const
 
std::vector< at::Tensormove_newstyle_outputs ()&&
 
template<>
NetDef GetSingleArgument (const std::string &name, const NetDef &default_value) const
 
template<>
vector< int > GetVectorFromIValueList (const c10::IValue &value) const
 
template<>
vector< float > GetVectorFromIValueList (const c10::IValue &value) const
 
template<>
vector< string > GetVectorFromIValueList (const c10::IValue &value) const
 
- Public Member Functions inherited from caffe2::Observable< OperatorBase >
 Observable (Observable &&)=default
 
Observableoperator= (Observable &&)=default
 
 C10_DISABLE_COPY_AND_ASSIGN (Observable)
 
const ObserverAttachObserver (std::unique_ptr< Observer > observer)
 
std::unique_ptr< ObserverDetachObserver (const Observer *observer_ptr)
 Returns a unique_ptr to the removed observer. More...
 
virtual size_t NumObservers ()
 
void StartAllObservers ()
 
void StopAllObservers ()
 

Data Fields

 USE_OPERATOR_CONTEXT_FUNCTIONS
 

Static Public Attributes

static constexpr int kSelfInputs = SparseFused ? 2 : 1
 
static constexpr int kNumInputs = Reducer::kInputCount + kSelfInputs
 
- Static Public Attributes inherited from caffe2::OperatorBase
static const int kNoNetPositionSet = -1
 

Additional Inherited Members

- Protected Member Functions inherited from caffe2::OperatorBase
virtual void RecordEvent (const char *=nullptr)
 
void SetEventFinished (const char *err_msg=nullptr)
 
void SetEventFinishedWithException (const char *err_msg=nullptr)
 
std::string getErrorMsg ()
 
 C10_DISABLE_COPY_AND_ASSIGN (OperatorBase)
 
- Protected Attributes inherited from caffe2::OperatorBase
std::unique_ptr< Eventevent_
 
- Protected Attributes inherited from caffe2::Observable< OperatorBase >
std::vector< std::unique_ptr< Observer > > observers_list_
 

Detailed Description

template<typename T, typename SIndex, class Context, class Reducer, bool SparseFused = true, class InputAccessor = BaseInputAccessor<T>>
class caffe2::AbstractUnsortedSegmentOp< T, SIndex, Context, Reducer, SparseFused, InputAccessor >

Unsorted segment reduction op with optional fused embedding lookup.

Base implementation for UnsortedSegmentXXX and UnsparseSortedSegmentXXX depending on SparseFused static argument.

Unlike the sorted version it allows to have "gaps" in segment ids.

Inputs: 0: DATA - input embedding to do lookups in 1..P: AUX_ARG_ - optional additional arguments to be passed to the reducer, should have the same first dimension as SEGMENT_IDS (e.g. scalars in WeightedSum)

if SparseFused == true:

P+1: INDICES - 1-D vector with indices to look up in DATA. Should have the same dimension as SEGMENT_IDS

P+1 if SparseFused == false:

P+1 or P+2: SEGMENT_IDS - unsorted segment ids 1-D vector

Args: num_segments - allows to override the dimension of the output. If not set it would be inferred from segment_ids tensor.

Output: Tensor with first dimension of K, where K is the max segment id + 1. Rest of dimensions are decided by reducer but usually are the same size as extra dimensions of DATA

Definition at line 987 of file segment_reduction_op.h.


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