Skip to content

Commit c0a6b5b

Browse files
committed
[ Layer ] Added Custom Multi Head Attention layer from FELICE
Incorporated the CPU version of custom MHA layer from FELICE repo Signed-off-by: Debadri Samaddar <s.debadri@samsung.com>
1 parent 9c8e874 commit c0a6b5b

File tree

8 files changed

+1289
-1
lines changed

8 files changed

+1289
-1
lines changed

api/ccapi/include/layer.h

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,9 @@ enum LayerType {
7575
ML_TRAIN_LAYER_TYPE_POSITIONAL_ENCODING, /**< Positional Encoding Layer type
7676
*/
7777
LAYER_IDENTITY = ML_TRAIN_LAYER_TYPE_IDENTITY, /**< Identity Layer type */
78+
LAYER_CUSTOM_MULTI_HEAD_ATTENTION =
79+
ML_TRAIN_LAYER_TYPE_CUSTOM_MULTI_HEAD_ATTENTION, /**< Multi Head Attention
80+
Layer type */
7881
LAYER_PREPROCESS_FLIP =
7982
ML_TRAIN_LAYER_TYPE_PREPROCESS_FLIP, /**< Preprocess flip Layer type */
8083
LAYER_PREPROCESS_TRANSLATE =
@@ -503,6 +506,16 @@ MultiHeadAttention(const std::vector<std::string> &properties = {}) {
503506
return createLayer(LayerType::LAYER_MULTI_HEAD_ATTENTION, properties);
504507
}
505508

509+
/**
510+
* @brief Helper function to create Custom Multi Head Attention Layer
511+
*/
512+
inline std::unique_ptr<Layer> CustomMultiHeadAttention(
513+
const std::vector<std::string> &properties = {},
514+
const LayerComputeEngine &compute_engine = LayerComputeEngine::CPU) {
515+
return createLayer(LayerType::LAYER_CUSTOM_MULTI_HEAD_ATTENTION, properties,
516+
compute_engine);
517+
}
518+
506519
/**
507520
* @brief Helper function to create Positional Encoding Layer
508521
*/

api/nntrainer-api-common.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,8 @@ typedef enum {
6363
ML_TRAIN_LAYER_TYPE_POSITIONAL_ENCODING =
6464
28, /**< Positional Encoding Layer type (Since 7.0) */
6565
ML_TRAIN_LAYER_TYPE_IDENTITY = 29, /**< Identity Layer type (Since 8.0) */
66+
ML_TRAIN_LAYER_TYPE_CUSTOM_MULTI_HEAD_ATTENTION =
67+
34, /**< Custom Multi Head Attention Layer type (Since 8.0) */
6668
ML_TRAIN_LAYER_TYPE_PREPROCESS_FLIP =
6769
300, /**< Preprocess flip Layer (Since 6.5) */
6870
ML_TRAIN_LAYER_TYPE_PREPROCESS_TRANSLATE =

nntrainer/app_context.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
#include <conv2d_layer.h>
4242
#include <cross_entropy_sigmoid_loss_layer.h>
4343
#include <cross_entropy_softmax_loss_layer.h>
44+
#include <custom_multi_head_attention_layer.h>
4445
#include <dropout.h>
4546
#include <embedding.h>
4647
#include <fc_layer.h>
@@ -299,6 +300,9 @@ static void add_default_object(AppContext &ac) {
299300
ac.registerFactory(nntrainer::createLayer<MultiHeadAttentionLayer>,
300301
MultiHeadAttentionLayer::type,
301302
LayerType::LAYER_MULTI_HEAD_ATTENTION);
303+
ac.registerFactory(nntrainer::createLayer<CustomMultiHeadAttentionLayer>,
304+
CustomMultiHeadAttentionLayer::type,
305+
LayerType::LAYER_CUSTOM_MULTI_HEAD_ATTENTION);
302306
ac.registerFactory(nntrainer::createLayer<ReduceMeanLayer>,
303307
ReduceMeanLayer::type, LayerType::LAYER_REDUCE_MEAN);
304308
ac.registerFactory(nntrainer::createLayer<PositionalEncodingLayer>,

0 commit comments

Comments
 (0)