116 lines
3.1 KiB
Protocol Buffer
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;
|
|
}
|