259 lines
9.8 KiB
C++
259 lines
9.8 KiB
C++
/*******************************************************************************
|
|
* Copyright (c) 2015-2018 Skymind, Inc.
|
|
*
|
|
* This program and the accompanying materials are made available under the
|
|
* terms of the Apache License, Version 2.0 which is available at
|
|
* https://www.apache.org/licenses/LICENSE-2.0.
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
|
|
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
|
|
* License for the specific language governing permissions and limitations
|
|
* under the License.
|
|
*
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
******************************************************************************/
|
|
|
|
//
|
|
// @author raver119@gmail.com
|
|
//
|
|
|
|
#ifndef LIBND4J_HEADERS_NN_H
|
|
#define LIBND4J_HEADERS_NN_H
|
|
|
|
#include <ops/declarable/headers/common.h>
|
|
|
|
namespace sd {
|
|
namespace ops {
|
|
|
|
#if NOT_EXCLUDED(OP_softmax)
|
|
DECLARE_CONFIGURABLE_OP(softmax, 1, 1, true, 0, 0);
|
|
DECLARE_CONFIGURABLE_OP(softmax_bp, 2, 1, true, 0, 0);
|
|
#endif
|
|
|
|
/**
|
|
* Local response normalization implementation as TF.
|
|
* input: 4D array
|
|
*
|
|
* T args:
|
|
*
|
|
* 0: bias
|
|
* 1: alpha
|
|
* 2: beta
|
|
*
|
|
* Int arg: depth - optional local radius
|
|
*
|
|
* output - 4D array
|
|
*/
|
|
#if NOT_EXCLUDED(OP_lrn)
|
|
DECLARE_CONFIGURABLE_OP(lrn, 1, 1, true, 3, 0);
|
|
#endif
|
|
|
|
/**
|
|
* Local response normalization - backprop variant.
|
|
* input:
|
|
* 0 - 4D array of data
|
|
* 1 - epsilon - 4D array of approximation
|
|
*
|
|
* T args:
|
|
*
|
|
* 0: bias
|
|
* 1: alpha
|
|
* 2: beta
|
|
*
|
|
* Int arg: depth - optional local radius
|
|
*
|
|
* output - next approximation as 4D array
|
|
*/
|
|
#if NOT_EXCLUDED(OP_lrn)
|
|
DECLARE_CONFIGURABLE_OP(lrn_bp, 2, 1, true, 3, 0);
|
|
#endif
|
|
|
|
/**
|
|
* Batch normalization implementation.
|
|
* Reference: https://arxiv.org/abs/1502.03167v3
|
|
*
|
|
* Expected arguments:
|
|
* input: input array (any number of dimensions)
|
|
* mean:
|
|
* variance:
|
|
* gamma:
|
|
* beta:
|
|
*
|
|
* Int args:
|
|
* 0: apply scale
|
|
* 1: apply offset
|
|
*
|
|
*
|
|
* T args:
|
|
* 0: epsilon
|
|
*/
|
|
#if NOT_EXCLUDED(OP_batchnorm)
|
|
DECLARE_CUSTOM_OP(batchnorm, 3, 1, false, 1, 2);
|
|
#endif
|
|
|
|
/**
|
|
* back prop in batch normalization
|
|
*
|
|
* Expected arguments:
|
|
* input: input array (any number of dimensions)
|
|
* mean:
|
|
* variance:
|
|
* gamma: optional
|
|
* beta: optional
|
|
* dLdOut: next epsilon
|
|
*
|
|
* Int args:
|
|
* 0: apply scale
|
|
* 1: apply offset
|
|
*
|
|
* T args:
|
|
* 0: epsilon
|
|
*
|
|
* output arrays:
|
|
* dL/dInput
|
|
* dL/dMean
|
|
* dL/dVariance
|
|
* dL/dGamma, optional
|
|
* dL/dBeta, optional
|
|
*/
|
|
#if NOT_EXCLUDED(OP_batchnorm)
|
|
DECLARE_CUSTOM_OP(batchnorm_bp, 4, 3, false, 1, 2);
|
|
#endif
|
|
|
|
|
|
/**
|
|
* This operation updates parameters with provided gradients, wrt learning rate
|
|
* Expected arguments:
|
|
* x: parameters, any shape
|
|
* y: gradients. same shape as x
|
|
* lr: optional, learning rate
|
|
*
|
|
* T args:
|
|
* 0: optional, learning rate
|
|
*/
|
|
#if NOT_EXCLUDED(OP_apply_sgd)
|
|
DECLARE_CONFIGURABLE_OP(apply_sgd, 2, 1, true, -2, 0);
|
|
#endif
|
|
|
|
/**
|
|
* This operation performs batch normalization of layer, it is based on following article https://arxiv.org/abs/1502.03167.
|
|
* Expected arguments:
|
|
* x: input 4D array of shape [bS,iH,iW,iD] (data format = NHWC) or [bS,iD,iH,iW] (data format = NCHW), where
|
|
* bS - batch size
|
|
* iH - input height
|
|
* iW - input width
|
|
* iD - input depth (or number of channels)
|
|
* scale: 1D input array of scale factors, shape [iD]
|
|
* offset: 1D input array of offsets (shifts), shape [iD]
|
|
* mean: 1D input array of population mean used for inference, shape [iD], this array is required only if isTraining = false
|
|
* variance: 1D input array of population mean used for inference, shape [iD], this array is required only if isTraining = false
|
|
*
|
|
* T input arguments:
|
|
* 0: epsilon, it is optional argument, default value is 0.001, this is small number to be added to the variance of x
|
|
*
|
|
* integer input arguments:
|
|
* 0: dataFormat, may have two values: zero -> NHWC, unity -> NCHW
|
|
* 1: isTraining, may have two values: zero -> inference, unity -> training
|
|
*/
|
|
#if NOT_EXCLUDED(OP_fused_batch_norm)
|
|
DECLARE_CUSTOM_OP(fused_batch_norm, 3, 1, false, 0, 2);
|
|
#endif
|
|
|
|
#if NOT_EXCLUDED(OP_log_softmax)
|
|
DECLARE_CONFIGURABLE_OP(log_softmax, 1, 1, true, 0, 0);
|
|
DECLARE_CONFIGURABLE_OP(log_softmax_bp, 2, 1, true, 0, 0);
|
|
#endif
|
|
|
|
|
|
/**
|
|
* relu_layer = relu(x*w + b)
|
|
*/
|
|
DECLARE_CUSTOM_OP(relu_layer, 3, 1, false, 0, 0);
|
|
|
|
/**
|
|
* applies layer normalization to input
|
|
* y = g * standardize(x) + b
|
|
*
|
|
* see sd::ops::standardize
|
|
*
|
|
*/
|
|
#if NOT_EXCLUDED(OP_layer_norm)
|
|
DECLARE_CONFIGURABLE_OP(layer_norm, 3, 1, true, 0, -2);
|
|
DECLARE_CUSTOM_OP(layer_norm_bp, 4, 1, false, 0, -2);
|
|
#endif
|
|
|
|
/**
|
|
* This operation performs dot product attention on the given timeseries input with the given queries
|
|
* out = sum(similarity(k_i, q) * v_i)
|
|
*
|
|
* similarity(k, q) = softmax(k * q) where x * q is the dot product of x and q
|
|
*
|
|
* Optionally with normalization step:
|
|
* similarity(k, q) = softmax(k * q / sqrt(size(q))
|
|
*
|
|
* See also "Attention is all you need" (https://arxiv.org/abs/1706.03762, p. 4, eq. 1)
|
|
*
|
|
* Note: This supports multiple queries at once, if only one query is available the queries vector still has to
|
|
* be 3D but can have queryCount = 1
|
|
*
|
|
* Note: keys and values usually is the same array. If you want to use it as the same array, simply pass it for
|
|
* both.
|
|
*
|
|
* Expected arguments:
|
|
* q: input 3D array "queries" of shape [batchSize, featureKeys, queryCount] or 4D array of shape [batchSize, numHeads, featureKeys, queryCount]
|
|
* k: input 3D array "keys" of shape [batchSize, featureKeys, timesteps] or 4D array of shape [batchSize, numHeads, featureKeys, timesteps]
|
|
* v: input 3D array "values" of shape [batchSize, featureValues, timesteps] or 4D array of shape [batchSize, numHeads, featureValues, timesteps]
|
|
* mask: OPTIONAL; array that defines which values should be skipped of shape [batchSize, timesteps]
|
|
*
|
|
* integer input arguments:
|
|
* 0: normalization, may have two values: zero -> do not apply normalization, one -> apply normalization
|
|
* 1: withWeights, may have two values: zero -> do not return weights, one -> return weights
|
|
*
|
|
* Output Arrays:
|
|
* 0: Attention result arrays of shape [batchSize, featureValues, queryCount] or [batchSize, numHeads, featureValues, queryCount]
|
|
* 1: OPTIONAL; Attention weights of shape [batchSize, timesteps, queryCount] or [batchSize, numHeads, timesteps, queryCount]
|
|
*/
|
|
#if NOT_EXCLUDED(OP_dot_product_attention)
|
|
DECLARE_CUSTOM_OP(dot_product_attention, 3, -1, false, 0, 2);
|
|
DECLARE_CUSTOM_OP(dot_product_attention_bp, 4, 3, false, 0, 1);
|
|
#endif
|
|
|
|
|
|
/**
|
|
* This performs multi-headed dot product attention on the given timeseries input
|
|
* out = concat(head_1, head_2, ..., head_n) * Wo
|
|
* head_i = dot_product_attention(Wq_i*q, Wk_i*k, Wv_i*v)
|
|
*
|
|
* Optionally with normalization when calculating the attention for each head.
|
|
*
|
|
* See also "Attention is all you need" (https://arxiv.org/abs/1706.03762, pp. 4,5, "3.2.2 Multi-Head Attention")
|
|
*
|
|
* This makes use of dot_product_attention OP support for rank 4 inputs.
|
|
*
|
|
* Expected arguments:
|
|
* q: input 3D array "queries" of shape [batchSize, featureKeys, queryCount]
|
|
* k: input 3D array "keys" of shape [batchSize, featureKeys, timesteps]
|
|
* v: input 3D array "values" of shape [batchSize, featureValues, timesteps]
|
|
* Wq: input query projection weights of shape [numHeads, projectedKeys, featureKeys]
|
|
* Wk: input key projection weights of shape [numHeads, projectedKeys, featureKeys]
|
|
* Wv: input value projection weights of shape [numHeads, projectedValues, featureValues]
|
|
* Wo: output projection weights of shape [numHeads * projectedValues, outSize]
|
|
* mask: OPTIONAL; array that defines which values should be skipped of shape [batchSize, timesteps]
|
|
*
|
|
* integer input arguments:
|
|
* 0: normalization, may have two values: zero -> do not apply normalization, one -> apply normalization
|
|
* 1: withWeights, may have two values: zero -> do not return weights, one -> return weights
|
|
*
|
|
* Output Arrays:
|
|
* 0: Attention result arrays of shape [batchSize, outSize, queryCount]
|
|
* 1: OPTIONAL; Attention weights of shape [batchSize, numHeads, timesteps, queryCount]
|
|
*/
|
|
#if NOT_EXCLUDED(OP_multi_head_dot_product_attention)
|
|
DECLARE_CUSTOM_OP(multi_head_dot_product_attention, 7, -1, false, 0, 2);
|
|
DECLARE_CUSTOM_OP(multi_head_dot_product_attention_bp, 8, 7, false, 0, 1);
|
|
#endif
|
|
}
|
|
}
|
|
|
|
#endif |