STT-tensorflow/tensorflow/stream_executor/dnn.proto

116 lines
3.1 KiB
Protocol Buffer

// LINT: LEGACY_NAMES
syntax = "proto3";
package stream_executor.dnn;
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/stream_executor";
// Specifies the data type used by an operation.
enum DataType {
kFloat = 0;
kDouble = 1;
kHalf = 2;
kInt8 = 3;
kInt32 = 4;
}
// Describes how a convolution input or output layer's data is formatted.
enum DataLayout {
// Naming convention:
// Y <-> row or height
// X <-> column or width
// Batch <-> batch, or N
// Depth <-> feature, or channel
// TODO(timshen): turn them into cuDNN names, e.g. kNCHW.
kYXDepthBatch = 0;
kYXBatchDepth = 1;
kBatchYXDepth = 2; // cuDNN's NHWC layout
kBatchDepthYX = 3; // cuDNN's NCHW layout
kBatchDepthYX4 = 4; // cuDNN's NCHW_VECT_C layout
}
// Describes how a convolution filter is laid out in the memory.
enum FilterLayout {
// Naming convention:
// Y <-> row or height
// X <-> column or width
// Output <-> output feature, or N
// Input <-> input feature, or N
// TODO(timshen): turn them into cuDNN names, e.g. kNCHW.
kOutputInputYX = 0; // cuDNN's NCHW layout
kOutputYXInput = 1; // cuDNN's NHWC layout
kOutputInputYX4 = 2; // cuDNN's NCHW_VECT_C layout
kInputYXOutput = 3;
kYXInputOutput = 4;
}
// Describes a kind of non-linearity (threshold-like mathematical function).
enum ActivationMode {
kNone = 0;
kSigmoid = 1;
// Rectified linear activation: f(x) = x < 0 ? 0 : x
kRelu = 2;
// Rectified linear activation; where upper maximum is 6.0.
kRelu6 = 3;
// Rectified linear activation; where upper maximum specified by
// BatchDescriptor::value_max().
kReluX = 4;
kTanh = 5;
// Like ReluX; but passes all values in the range [-X,X].
kBandPass = 6;
}
// Describe the math definition for the conv op. The popular behavior is
// actually called cross-correlation in math, despite the operation is often
// referred as convolution. See cuDNN cudnnConvolutionMode_t.
enum ConvolutionMode {
CROSS_CORRELATION = 0;
CONVOLUTION = 1;
}
enum ConvolutionKind {
INVALID = 0;
FORWARD = 1;
BACKWARD_FILTER = 2;
BACKWARD_DATA = 3;
FORWARD_BIAS_ACTIVATION = 4;
}
// Generic tensor representation.
message TensorDescriptorProto {
repeated int64 dimensions = 1;
DataType data_type = 2;
oneof layout_oneof {
DataLayout data_layout = 3;
FilterLayout filter_layout = 4;
}
}
// Generic algorithm representation.
message AlgorithmProto {
enum MathType {
DEFAULT_MATH = 0;
// The GPU may operate 4x4 matrix FMA.
// See cuDNN's documentation for CUDNN_TENSOR_OP_MATH.
TENSOR_OP_MATH = 1;
}
int64 algo_id = 1;
MathType math_type = 2;
}
// Convolution-specific parameters.
message ConvolutionDescriptorProto {
repeated int64 paddings = 1;
repeated int64 strides = 2;
repeated int64 dilations = 3;
// The "accumulator" type. For example, use F32 as an accumulator for F16
// convolutions.
// See cuDNN's cudnnConvolutionMode_t.
DataType compute_mode = 4;
// See cuDNN's group count.
int32 group_count = 5;
ConvolutionMode convolution_mode = 6;
// Tensorflow node name, same as in NodeDef, for debugging purposes.
string name = 7;
}