From 09e299534615628b02e79aa92fe18ce022bb684a Mon Sep 17 00:00:00 2001 From: Mrityunjay Tripathi Date: Wed, 13 Jan 2021 09:19:09 +0530 Subject: [PATCH] update multihead attention layer --- src/mlpack/methods/ann/layer/layer_types.hpp | 2 +- .../methods/ann/layer/multihead_attention.hpp | 106 +++--- .../ann/layer/multihead_attention_impl.hpp | 109 +++--- src/mlpack/tests/ann_layer_test.cpp | 334 +++++++++--------- 4 files changed, 276 insertions(+), 275 deletions(-) diff --git a/src/mlpack/methods/ann/layer/layer_types.hpp b/src/mlpack/methods/ann/layer/layer_types.hpp index f2f0d1aff8..7f9872b204 100644 --- a/src/mlpack/methods/ann/layer/layer_types.hpp +++ b/src/mlpack/methods/ann/layer/layer_types.hpp @@ -33,7 +33,7 @@ #include #include #include -// #include +#include // #include // #include // #include diff --git a/src/mlpack/methods/ann/layer/multihead_attention.hpp b/src/mlpack/methods/ann/layer/multihead_attention.hpp index 3421fa4183..3160c22994 100644 --- a/src/mlpack/methods/ann/layer/multihead_attention.hpp +++ b/src/mlpack/methods/ann/layer/multihead_attention.hpp @@ -48,24 +48,24 @@ namespace ann /** Artificial Neural Network. */ { * of shape `(embedDim * tgtSeqLen, batchSize)`. The embeddings are stored * consequently. * - * @tparam InputDataType Type of the input data (arma::colvec, arma::mat, + * @tparam InputType Type of the input data (arma::colvec, arma::mat, * arma::sp_mat or arma::cube). - * @tparam OutputDataType Type of the output data (arma::colvec, arma::mat, + * @tparam OutputType Type of the output data (arma::colvec, arma::mat, * arma::sp_mat or arma::cube). * @tparam RegularizerType Type of the regularizer to be used. */ template < - typename InputDataType = arma::mat, - typename OutputDataType = arma::mat, + typename InputType = arma::mat, + typename OutputType = arma::mat, typename RegularizerType = NoRegularizer > -class MultiheadAttention +class MultiheadAttentionType : public Layer { public: /** * Default constructor. */ - MultiheadAttention(); + MultiheadAttentionType(); /** * Create the MultiheadAttention object using the specified modules. @@ -74,11 +74,15 @@ class MultiheadAttention * @param srcSeqLen Source sequence length. * @param embedDim Total dimension of the model. * @param numHeads Number of parallel attention heads. + * @param attnMask Two dimensional Attention Mask. + * @param keyPaddingMask Key Padding Mask. */ - MultiheadAttention(const size_t tgtSeqLen, - const size_t srcSeqLen, - const size_t embedDim, - const size_t numHeads); + MultiheadAttentionType(const size_t tgtSeqLen, + const size_t srcSeqLen, + const size_t embedDim, + const size_t numHeads, + const InputType& attnmask = InputType(), + const InputType& keyPaddingMask = InputType()); /** * Reset the layer parameters. @@ -92,8 +96,7 @@ class MultiheadAttention * @param input The query matrix. * @param output Resulting output activation. */ - template - void Forward(const arma::Mat& input, arma::Mat& output); + void Forward(const InputType& input, OutputType& output); /** * Ordinary feed backward pass of a neural network, calculating the function @@ -103,10 +106,9 @@ class MultiheadAttention * @param gy The backpropagated error. * @param g The calculated gradient. */ - template - void Backward(const arma::Mat& /* input */, - const arma::Mat& gy, - arma::Mat& g); + void Backward(const InputType& /* input */, + const OutputType& gy, + OutputType& g); /** * Calculate the gradient using the output delta and the input activation. @@ -115,10 +117,9 @@ class MultiheadAttention * @param error The calculated error. * @param gradient The calculated gradient. */ - template - void Gradient(const arma::Mat& input, - const arma::Mat& error, - arma::Mat& gradient); + void Gradient(const InputType& input, + const OutputType& error, + OutputType& gradient); /** * Serialize the layer. @@ -147,46 +148,46 @@ class MultiheadAttention size_t& NumHeads() { return numHeads; } //! Get the two dimensional Attention Mask. - OutputDataType const& AttentionMask() const { return attnMask; } + OutputType const& AttentionMask() const { return attnMask; } //! Modify the two dimensional Attention Mask. - OutputDataType& AttentionMask() { return attnMask; } + OutputType& AttentionMask() { return attnMask; } //! Get Key Padding Mask. - OutputDataType const& KeyPaddingMask() const { return keyPaddingMask; } + OutputType const& KeyPaddingMask() const { return keyPaddingMask; } //! Modify the Key Padding Mask. - OutputDataType& KeyPaddingMask() { return keyPaddingMask; } + OutputType& KeyPaddingMask() { return keyPaddingMask; } //! Get the output parameter. - OutputDataType const& OutputParameter() const { return outputParameter; } + OutputType const& OutputParameter() const { return outputParameter; } //! Modify the output parameter. - OutputDataType& OutputParameter() { return outputParameter; } + OutputType& OutputParameter() { return outputParameter; } //! Get the delta. - OutputDataType const& Delta() const { return delta; } + OutputType const& Delta() const { return delta; } //! Modify the delta. - OutputDataType& Delta() { return delta; } + OutputType& Delta() { return delta; } //! Get the gradient. - OutputDataType const& Gradient() const { return grad; } + OutputType const& Gradient() const { return grad; } //! Modify the gradient. - OutputDataType& Gradient() { return grad; } + OutputType& Gradient() { return grad; } //! Get the parameters. - OutputDataType const& Parameters() const { return weights; } + OutputType const& Parameters() const { return weights; } //! Modify the parameters. - OutputDataType& Parameters() { return weights; } + OutputType& Parameters() { return weights; } private: //! Element Type of the input. - typedef typename OutputDataType::elem_type ElemType; + typedef typename OutputType::elem_type ElemType; //! Target sequence length. size_t tgtSeqLen; - //! Source sequence lenght. + //! Source sequence length. size_t srcSeqLen; - //! Locally-stored module output size. + //! Locally-stored dimensionality of each embedding vector. size_t embedDim; //! Locally-stored number of parallel attention heads. @@ -196,37 +197,37 @@ class MultiheadAttention size_t headDim; //! Two dimensional Attention Mask of shape (tgtSeqLen, srcSeqLen). - OutputDataType attnMask; + OutputType attnMask; //! Key Padding Mask. - OutputDataType keyPaddingMask; + OutputType keyPaddingMask; //! Locally-stored weight matrix associated with query. - OutputDataType queryWt; + OutputType queryWt; //! Locally-stored weight matrix associated with key. - OutputDataType keyWt; + OutputType keyWt; //! Locally-stored weight matrix associated with value. - OutputDataType valueWt; + OutputType valueWt; //! Locally-stored weight matrix associated with attnWt. - OutputDataType outWt; + OutputType outWt; //! Locally-stored bias associated with query. - OutputDataType qBias; + OutputType qBias; //! Locally-stored bias associated with key. - OutputDataType kBias; + OutputType kBias; //! Locall-stored bias associated with value. - OutputDataType vBias; + OutputType vBias; //! Locally-stored bias associated with attnWt. - OutputDataType outBias; + OutputType outBias; //! Locally-stored weights parameter. - OutputDataType weights; + OutputType weights; //! Locally-stored projected query matrix over linear layer. arma::Cube qProj; @@ -244,20 +245,25 @@ class MultiheadAttention arma::Cube attnOut; //! Softmax layer to represent the probabilities of next sequence. - Softmax softmax; + Softmax softmax; //! Locally-stored delta object. - OutputDataType delta; + OutputType delta; //! Locally-stored gradient. - OutputDataType grad; + OutputType grad; //! Locally-stored output parameter. - OutputDataType outputParameter; + OutputType outputParameter; //! Locally-stored regularizer object. RegularizerType regularizer; }; // class MultiheadAttention + +// Standard MultiheadAttention layer using no regularization. +typedef MultiheadAttentionType +MultiheadAttention; + } // namespace ann } // namespace mlpack diff --git a/src/mlpack/methods/ann/layer/multihead_attention_impl.hpp b/src/mlpack/methods/ann/layer/multihead_attention_impl.hpp index d2da8788e9..fdd1367c99 100644 --- a/src/mlpack/methods/ann/layer/multihead_attention_impl.hpp +++ b/src/mlpack/methods/ann/layer/multihead_attention_impl.hpp @@ -21,31 +21,35 @@ namespace mlpack { namespace ann /** Artificial Neural Network. */ { -template -MultiheadAttention:: -MultiheadAttention() : +template +MultiheadAttentionType:: +MultiheadAttentionType() : tgtSeqLen(0), srcSeqLen(0), embedDim(0), numHeads(0), - headDim(0) + headDim(0), + attnMask(InputType()), + keyPaddingMask(InputType()) { // Nothing to do here. } -template -MultiheadAttention:: -MultiheadAttention( +template +MultiheadAttentionType:: +MultiheadAttentionType( const size_t tgtSeqLen, const size_t srcSeqLen, const size_t embedDim, - const size_t numHeads) : + const size_t numHeads, + const InputType& attnMask, + const InputType& keyPaddingMask) : tgtSeqLen(tgtSeqLen), srcSeqLen(srcSeqLen), embedDim(embedDim), - numHeads(numHeads) + numHeads(numHeads), + attnMask(attnMask), + keyPaddingMask(keyPaddingMask) { if (embedDim % numHeads != 0) { @@ -57,38 +61,33 @@ MultiheadAttention( weights.set_size(4 * (embedDim + 1) * embedDim, 1); } -template -void MultiheadAttention:: +template +void MultiheadAttentionType:: Reset() { - typedef typename arma::Mat MatType; - - queryWt = MatType(weights.memptr(), embedDim, embedDim, false, false); - keyWt = MatType(weights.memptr() + embedDim * embedDim, + queryWt = OutputType(weights.memptr(), embedDim, embedDim, false, false); + keyWt = OutputType(weights.memptr() + embedDim * embedDim, embedDim, embedDim, false, false); - valueWt = MatType(weights.memptr() + 2 * embedDim * embedDim, + valueWt = OutputType(weights.memptr() + 2 * embedDim * embedDim, embedDim, embedDim, false, false); - outWt = MatType(weights.memptr() + 3 * embedDim * embedDim, + outWt = OutputType(weights.memptr() + 3 * embedDim * embedDim, embedDim, embedDim, false, false); - qBias = MatType(weights.memptr() + qBias = OutputType(weights.memptr() + 4 * embedDim * embedDim, embedDim, 1, false, false); - kBias = MatType(weights.memptr() + kBias = OutputType(weights.memptr() + (4 * embedDim + 1) * embedDim, embedDim, 1, false, false); - vBias = MatType(weights.memptr() + vBias = OutputType(weights.memptr() + (4 * embedDim + 2) * embedDim, embedDim, 1, false, false); - outBias = MatType(weights.memptr() + outBias = OutputType(weights.memptr() + (4 * embedDim + 3) * embedDim, 1, embedDim, false, false); } -template -template -void MultiheadAttention:: -Forward(const arma::Mat& input, arma::Mat& output) +template +void MultiheadAttentionType:: +Forward(const InputType& input, OutputType& output) { - typedef typename arma::Cube CubeType; + typedef typename arma::Cube CubeType; if (input.n_rows != embedDim * (tgtSeqLen + 2 * srcSeqLen)) { @@ -104,12 +103,12 @@ Forward(const arma::Mat& input, arma::Mat& output) // The shape of q : (embedDim, tgtSeqLen, batchSize). // The shape of k : (embedDim, srcSeqLen, batchSize). // The shape of v : (embedDim, srcSeqLen, batchSize). - const CubeType q(const_cast&>(input).memptr(), + const CubeType q(const_cast(input).memptr(), embedDim, tgtSeqLen, batchSize, false, false); - const CubeType k(const_cast&>(input).memptr() + + const CubeType k(const_cast(input).memptr() + embedDim * tgtSeqLen * batchSize, embedDim, srcSeqLen, batchSize, false, false); - const CubeType v(const_cast&>(input).memptr() + + const CubeType v(const_cast(input).memptr() + embedDim * (tgtSeqLen + srcSeqLen) * batchSize, embedDim, srcSeqLen, batchSize, false, false); @@ -188,15 +187,13 @@ Forward(const arma::Mat& input, arma::Mat& output) } } -template -template -void MultiheadAttention:: -Backward(const arma::Mat& /* input */, - const arma::Mat& gy, - arma::Mat& g) +template +void MultiheadAttentionType:: +Backward(const InputType& /* input */, + const OutputType& gy, + OutputType& g) { - typedef typename arma::Cube CubeType; + typedef typename arma::Cube CubeType; if (gy.n_rows != tgtSeqLen * embedDim) { @@ -210,7 +207,7 @@ Backward(const arma::Mat& /* input */, // The shape of gyTemp : (tgtSeqLen, embedDim, batchSize). // We need not split it into n heads now because this is the part when // output were concatenated from n heads. - CubeType gyTemp(const_cast&>(gy).memptr(), embedDim, + CubeType gyTemp(const_cast(gy).memptr(), embedDim, tgtSeqLen, batchSize, true, false); // The shape of gyTemp : (embedDim, tgtSeqLen, batchSize). @@ -280,16 +277,13 @@ Backward(const arma::Mat& /* input */, } } -template -template -void MultiheadAttention:: -Gradient(const arma::Mat& input, - const arma::Mat& error, - arma::Mat& gradient) +template +void MultiheadAttentionType:: +Gradient(const InputType& input, + const OutputType& error, + OutputType& gradient) { - typedef typename arma::Cube CubeType; - typedef typename arma::Mat MatType; + typedef typename arma::Cube CubeType; if (input.n_rows != embedDim * (tgtSeqLen + 2 * srcSeqLen)) { @@ -307,16 +301,16 @@ Gradient(const arma::Mat& input, // The shape of gradient : (4 * embedDim * embedDim + 4 * embedDim, 1). gradient.set_size(arma::size(weights)); - const CubeType q(const_cast(input).memptr(), + const CubeType q(const_cast(input).memptr(), embedDim, tgtSeqLen, batchSize, false, false); - const CubeType k(const_cast(input).memptr() + q.n_elem, + const CubeType k(const_cast(input).memptr() + q.n_elem, embedDim, srcSeqLen, batchSize, false, false); - const CubeType v(const_cast(input).memptr() + q.n_elem + k.n_elem, + const CubeType v(const_cast(input).memptr() + q.n_elem + k.n_elem, embedDim, srcSeqLen, batchSize, false, false); // Reshape the propagated error into a cube. // The shape of errorTemp : (embedDim, tgtSeqLen, batchSize). - CubeType errorTemp(const_cast&>(error).memptr(), embedDim, + CubeType errorTemp(const_cast(error).memptr(), embedDim, tgtSeqLen, batchSize, true, false); // Gradient wrt. outBias, i.e. dL/d(outBias). @@ -430,10 +424,9 @@ Gradient(const arma::Mat& input, regularizer.Evaluate(weights, gradient); } -template +template template -void MultiheadAttention:: +void MultiheadAttentionType:: serialize(Archive& ar, const uint32_t /* version */) { ar(CEREAL_NVP(tgtSeqLen)); diff --git a/src/mlpack/tests/ann_layer_test.cpp b/src/mlpack/tests/ann_layer_test.cpp index 974f15663b..31645a983d 100644 --- a/src/mlpack/tests/ann_layer_test.cpp +++ b/src/mlpack/tests/ann_layer_test.cpp @@ -4650,199 +4650,201 @@ TEST_CASE("AdaptiveMeanPoolingTestCase", "[ANNLayerTest]") // } // } -// /** -// * Simple Multihead Attention test. -// */ -// TEST_CASE("SimpleMultiheadAttentionTest", "[ANNLayerTest]") -// { -// size_t tLen = 5; -// size_t sLen = tLen; -// size_t embedDim = 4; -// size_t numHeads = 2; -// size_t bsz = 3; +/** + * Simple Multihead Attention test. + */ +TEST_CASE("SimpleMultiheadAttentionTest", "[ANNLayerTest]") +{ + size_t tLen = 5; + size_t sLen = tLen; + size_t embedDim = 4; + size_t numHeads = 2; + size_t bsz = 3; -// arma::mat query = 0.1 * arma::randu(embedDim * tLen, bsz); -// arma::mat output; + arma::mat query = 0.1 * arma::randu(embedDim * tLen, bsz); + arma::mat output; -// arma::mat attnMask = arma::zeros(tLen, sLen); -// for (size_t i = 0; i < tLen; ++i) -// { -// for (size_t j = 0; j < sLen; ++j) -// { -// if (i < j) -// attnMask(i, j) = std::numeric_limits::lowest(); -// } -// } + arma::mat attnMask = arma::zeros(tLen, sLen); + for (size_t i = 0; i < tLen; ++i) + { + for (size_t j = 0; j < sLen; ++j) + { + if (i < j) + attnMask(i, j) = std::numeric_limits::lowest(); + } + } -// arma::mat keyPaddingMask = arma::zeros(1, sLen); -// keyPaddingMask(sLen - 1) = std::numeric_limits::lowest(); + arma::mat keyPaddingMask = arma::zeros(1, sLen); + keyPaddingMask(sLen - 1) = std::numeric_limits::lowest(); -// MultiheadAttention<> module(tLen, sLen, embedDim, numHeads); -// module.AttentionMask() = attnMask; -// module.KeyPaddingMask() = keyPaddingMask; -// module.Reset(); -// module.Parameters().randu(); + MultiheadAttention module(tLen, sLen, embedDim, numHeads); + module.AttentionMask() = attnMask; + module.KeyPaddingMask() = keyPaddingMask; + module.Reset(); + module.Parameters().randu(); -// // Forward test. -// arma::mat input = arma::join_cols(arma::join_cols(query, query), query); + // Forward test. + arma::mat input = arma::join_cols(arma::join_cols(query, query), query); -// module.Forward(input, output); -// REQUIRE(output.n_rows == embedDim * tLen); -// REQUIRE(output.n_cols == bsz); + module.Forward(input, output); + REQUIRE(output.n_rows == embedDim * tLen); + REQUIRE(output.n_cols == bsz); -// // Backward test. -// arma::mat gy = 0.01 * arma::randu(embedDim * tLen, bsz); -// arma::mat g; -// module.Backward(input, gy, g); -// REQUIRE(g.n_rows == input.n_rows); -// REQUIRE(g.n_cols == input.n_cols); + // Backward test. + arma::mat gy = 0.01 * arma::randu(embedDim * tLen, bsz); + arma::mat g; + module.Backward(input, gy, g); + REQUIRE(g.n_rows == input.n_rows); + REQUIRE(g.n_cols == input.n_cols); -// // Gradient test. -// arma::mat error = 0.05 * arma::randu(embedDim * tLen, bsz); -// arma::mat gradient; -// module.Gradient(input, error, gradient); -// REQUIRE(gradient.n_rows == module.Parameters().n_rows); -// REQUIRE(gradient.n_cols == module.Parameters().n_cols); -// } + // Gradient test. + arma::mat error = 0.05 * arma::randu(embedDim * tLen, bsz); + arma::mat gradient; + module.Gradient(input, error, gradient); + REQUIRE(gradient.n_rows == module.Parameters().n_rows); + REQUIRE(gradient.n_cols == module.Parameters().n_cols); +} -// /** -// * Jacobian MultiheadAttention module test. -// */ -// TEST_CASE("JacobianMultiheadAttentionTest", "[ANNLayerTest]") -// { -// // Check when query = key = value. -// for (size_t i = 0; i < 5; ++i) -// { -// const size_t tgtSeqLen = 2; -// const size_t embedDim = 4; -// const size_t nHeads = 2; -// const size_t batchSize = 1; +/** + * Jacobian MultiheadAttention module test. + */ +TEST_CASE("JacobianMultiheadAttentionTest", "[ANNLayerTest]") +{ + // Check when query = key = value. + for (size_t i = 0; i < 5; ++i) + { + const size_t tgtSeqLen = 2; + const size_t embedDim = 4; + const size_t nHeads = 2; + const size_t batchSize = 1; -// arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize); -// arma::mat input = arma::join_cols(arma::join_cols(query, query), query); + arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize); + arma::mat input = arma::join_cols(arma::join_cols(query, query), query); -// MultiheadAttention<> module(tgtSeqLen, tgtSeqLen, embedDim, nHeads); -// module.Parameters().randu(); + MultiheadAttention module(tgtSeqLen, tgtSeqLen, embedDim, nHeads); + module.Parameters().randu(); -// double error = CustomJacobianTest(module, input); -// REQUIRE(error <= 1e-5); -// } + double error = CustomJacobianTest(module, input); + REQUIRE(error <= 1e-5); + } -// // Check when key = value. -// for (size_t i = 0; i < 5; ++i) -// { -// const size_t tgtSeqLen = 2; -// const size_t srcSeqLen = math::RandInt(2, 5); -// const size_t embedDim = 4; -// const size_t nHeads = 2; -// const size_t batchSize = 1; + // Check when key = value. + for (size_t i = 0; i < 5; ++i) + { + const size_t tgtSeqLen = 2; + const size_t srcSeqLen = math::RandInt(2, 5); + const size_t embedDim = 4; + const size_t nHeads = 2; + const size_t batchSize = 1; -// arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize); -// arma::mat key = 0.091 * arma::randu(embedDim * srcSeqLen, batchSize); -// arma::mat input = arma::join_cols(arma::join_cols(query, key), key); + arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize); + arma::mat key = 0.091 * arma::randu(embedDim * srcSeqLen, batchSize); + arma::mat input = arma::join_cols(arma::join_cols(query, key), key); -// MultiheadAttention<> module(tgtSeqLen, srcSeqLen, embedDim, nHeads); -// module.Parameters().randu(); + MultiheadAttention module(tgtSeqLen, srcSeqLen, embedDim, nHeads); + module.Parameters().randu(); -// double error = CustomJacobianTest(module, input); -// REQUIRE(error <= 1e-5); -// } + double error = CustomJacobianTest(module, input); + REQUIRE(error <= 1e-5); + } -// // Check when query, key and value are not same. -// for (size_t i = 0; i < 5; ++i) -// { -// const size_t tgtSeqLen = 2; -// const size_t srcSeqLen = math::RandInt(2, 5); -// const size_t embedDim = 4; -// const size_t nHeads = 2; -// const size_t batchSize = 1; + // Check when query, key and value are not same. + for (size_t i = 0; i < 5; ++i) + { + const size_t tgtSeqLen = 2; + const size_t srcSeqLen = math::RandInt(2, 5); + const size_t embedDim = 4; + const size_t nHeads = 2; + const size_t batchSize = 1; -// arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize); -// arma::mat key = 0.091 * arma::randu(embedDim * srcSeqLen, batchSize); -// arma::mat value = 0.045 * arma::randu(embedDim * srcSeqLen, batchSize); -// arma::mat input = arma::join_cols(arma::join_cols(query, key), value); + arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize); + arma::mat key = 0.091 * arma::randu(embedDim * srcSeqLen, batchSize); + arma::mat value = 0.045 * arma::randu(embedDim * srcSeqLen, batchSize); + arma::mat input = arma::join_cols(arma::join_cols(query, key), value); -// MultiheadAttention<> module(tgtSeqLen, srcSeqLen, embedDim, nHeads); -// module.Parameters().randu(); + MultiheadAttention module(tgtSeqLen, srcSeqLen, embedDim, nHeads); + module.Parameters().randu(); -// double error = JacobianTest(module, input); -// REQUIRE(error <= 1e-5); -// } -// } + double error = JacobianTest(module, input); + REQUIRE(error <= 1e-5); + } +} -// /** -// * Numerical gradient test for MultiheadAttention layer. -// */ -// TEST_CASE("GradientMultiheadAttentionTest", "[ANNLayerTest]") -// { -// struct GradientFunction -// { -// GradientFunction() : -// tgtSeqLen(2), -// srcSeqLen(2), -// embedDim(4), -// nHeads(2), -// vocabSize(5), -// batchSize(2) -// { -// input = arma::randu(embedDim * (tgtSeqLen + 2 * srcSeqLen), batchSize); -// target = arma::zeros(vocabSize, batchSize); -// for (size_t i = 0; i < target.n_elem; ++i) -// { -// const size_t label = mlpack::math::RandInt(1, vocabSize); -// target(i) = label; -// } +/** + * Numerical gradient test for MultiheadAttention layer. + */ +TEST_CASE("GradientMultiheadAttentionTest", "[ANNLayerTest]") +{ + struct GradientFunction + { + GradientFunction() : + tgtSeqLen(2), + srcSeqLen(2), + embedDim(4), + nHeads(2), + vocabSize(5), + batchSize(2) + { + input = arma::randu(embedDim * (tgtSeqLen + 2 * srcSeqLen), batchSize); + target = arma::zeros(vocabSize, batchSize); + for (size_t i = 0; i < target.n_elem; ++i) + { + const size_t label = mlpack::math::RandInt(1, vocabSize); + target(i) = label; + } -// attnMask = arma::zeros(tgtSeqLen, srcSeqLen); -// for (size_t i = 0; i < tgtSeqLen; ++i) -// { -// for (size_t j = 0; j < srcSeqLen; ++j) -// { -// if (i < j) -// attnMask(i, j) = std::numeric_limits::lowest(); -// } -// } + attnMask = arma::zeros(tgtSeqLen, srcSeqLen); + for (size_t i = 0; i < tgtSeqLen; ++i) + { + for (size_t j = 0; j < srcSeqLen; ++j) + { + if (i < j) + attnMask(i, j) = std::numeric_limits::lowest(); + } + } -// keyPaddingMask = arma::zeros(1, srcSeqLen); -// keyPaddingMask(srcSeqLen - 1) = std::numeric_limits::lowest(); + keyPaddingMask = arma::zeros(1, srcSeqLen); + keyPaddingMask(srcSeqLen - 1) = std::numeric_limits::lowest(); -// model = new FFN, XavierInitialization>(); -// model->Predictors() = input; -// model->Responses() = target; -// attnModule = new MultiheadAttention<>(tgtSeqLen, srcSeqLen, -// embedDim, nHeads); -// attnModule->AttentionMask() = attnMask; -// attnModule->KeyPaddingMask() = keyPaddingMask; -// model->Add(attnModule); -// model->Add>(embedDim * tgtSeqLen, vocabSize); -// model->Add>(); -// } + model = new FFN, XavierInitialization>(); + model->Predictors() = input; + model->Responses() = target; + // attnModule = new MultiheadAttention(tgtSeqLen, srcSeqLen, embedDim, + // nHeads); + // attnModule->AttentionMask() = attnMask; + // attnModule->KeyPaddingMask() = keyPaddingMask; + // model->Add(attnModule); + model->Add(tgtSeqLen, srcSeqLen, embedDim, nHeads, + attnMask, keyPaddingMask); + model->Add(embedDim * tgtSeqLen, vocabSize); + model->Add(); + } -// ~GradientFunction() -// { -// delete model; -// } + ~GradientFunction() + { + delete model; + } -// double Gradient(arma::mat& gradient) const -// { -// double error = model->Evaluate(model->Parameters(), 0, batchSize); -// model->Gradient(model->Parameters(), 0, gradient, batchSize); -// return error; -// } + double Gradient(arma::mat& gradient) const + { + double error = model->Evaluate(model->Parameters(), 0, batchSize); + model->Gradient(model->Parameters(), 0, gradient, batchSize); + return error; + } -// arma::mat& Parameters() { return model->Parameters(); } + arma::mat& Parameters() { return model->Parameters(); } -// FFN, XavierInitialization>* model; -// MultiheadAttention<>* attnModule; + FFN, XavierInitialization>* model; + // MultiheadAttention* attnModule; -// arma::mat input, target, attnMask, keyPaddingMask; -// const size_t tgtSeqLen; -// const size_t srcSeqLen; -// const size_t embedDim; -// const size_t nHeads; -// const size_t vocabSize; -// const size_t batchSize; -// } function; + arma::mat input, target, attnMask, keyPaddingMask; + const size_t tgtSeqLen; + const size_t srcSeqLen; + const size_t embedDim; + const size_t nHeads; + const size_t vocabSize; + const size_t batchSize; + } function; -// REQUIRE(CheckGradient(function) <= 3e-06); -// } + REQUIRE(CheckGradient(function) <= 3e-06); +}