wehub/deeplearning4j--deeplearning4j
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
4e5a78dcb79d84c2ea88f9a823da09dca5a65a14
deeplearning4j--deeplearning4j/libnd4j/include/ops/declarable/helpers/cpu/ctcLoss.cpp
T
wehub-resource-sync 4f3b7da785 chore: import upstream snapshot with attribution
2026-07-13 12:47:05 +08:00

433 lines
17 KiB
C++
Raw Blame History

/*******************************************************************************
* Copyright (c) 2021 Deeplearning4j Contributors
*
* 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 AbdelRauf
//
#include <execution/ThreadPool.h>
#include <execution/Threads.h>
#include <helpers/LoopsCoordsHelper.h>
#include <ops/declarable/helpers/ctc.h>
#include <cmath>
#include <memory>
#include <stdexcept>
#include <type_traits>
#include <system/selective_rendering.h>
#if NOT_EXCLUDED(OP_ctc_loss)
namespace sd {
namespace ops {
namespace helpers {
template <bool IsLogPStrided = false, bool IsLblStrided = false, typename Type, typename IndexType>
Type forward(Type *alphaPtr, const sd::LongType &incA, const Type *logP, const sd::LongType &incP, const IndexType *lbl,
const sd::LongType &lenSB, const sd::LongType &lenT, const int &blankIndex, int elwiseP = 1,
int elwiseS = 1) {
Type negInf = negative_infinity<Type>();
// initialize alphas at t=0
alphaPtr[0] = element<IsLogPStrided>(logP, blankIndex, elwiseP);
// alphaPtr[1] =logP[lbl[0]];
alphaPtr[1] = element<IsLogPStrided>(logP, *lbl, elwiseP);
// the rest initialization was skipped
// as its assumed the array already were initialized with negative infinity
// move to the next frame
Type *alphaPrevPtr = alphaPtr;
alphaPtr += incA;
logP += incP;
auto startX = lenSB - 2 * lenT;
// process the rest
for (auto t = 1; t < lenT; t++) {
// start = max(0,L-2*(T-t))
auto s = startX + 2 * t;
s = s > 0 ? s : 0;
for (; s < lenSB; s++) {
auto ind = s / 2; // our real index
// we force blanks for even indexes
// strided version of lbl[ind] => element<IsLblStrided>(lbl, ind, elwiseS)
auto currentInd = (s % 2 == 0) ? blankIndex : element<IsLblStrided>(lbl, ind, elwiseS);
// {t-1,s}
Type alphaS = alphaPrevPtr[s];
Type alphaS_1 = s > 0 ? alphaPrevPtr[s - 1] : negInf;
// logP[currentInd] or logP[currentInd*elwiseP]
auto currentProb = element<IsLogPStrided>(logP, currentInd, elwiseP);
// if blank or the same as previous
if (s > 1 && currentInd != blankIndex && currentInd != element<IsLblStrided>(lbl, ind - 1, elwiseS)) {
Type alphaS_2 = alphaPrevPtr[s - 2];
alphaPtr[s] = log_sum_exp(alphaS, alphaS_1, alphaS_2) + currentProb;
} else {
alphaPtr[s] = log_sum_exp(alphaS, alphaS_1) + currentProb;
}
}
// store t-1 alpha Ptr
alphaPrevPtr = alphaPtr;
logP += incP;
alphaPtr += incA;
}
auto logP0 = alphaPrevPtr[lenSB - 1];
auto logP1 = alphaPrevPtr[lenSB - 2];
return -log_sum_exp(logP0, logP1);
}
//#undef CALCULATE_ALL_IN_ONE_FRAME_LOOP
template <bool IsLogPStrided = false, bool IsLblStrided = false, bool isGradStrided = false, typename Type,
typename IndexType = int>
void backwardAndGrad(Type forwardLogLoss, Type *alphaPtr, Type *bettaPtr, int incA, const Type *logP, int incP,
Type *gradPtr, int incG, const IndexType *lbl, const sd::LongType &lenS, const sd::LongType &lenT,
const sd::LongType &lenK, const int &blankIndex, int elwiseP = 1, int elwiseS = 1,
int elwiseG = 1) {
Type negInf = negative_infinity<Type>();
sd::LongType lenSB = 2 * lenS + 1;
auto origBetta = bettaPtr;
auto origLogP = logP;
// move to the last frame
bettaPtr += (lenT - 1) * incA;
logP += (lenT - 1) * incP;
// initialize bettas at t=lenT
bettaPtr[lenSB - 1] = element<IsLogPStrided>(logP, blankIndex, elwiseP);
auto lblIndex = element<IsLblStrided>(lbl, lenS - 1, elwiseS);
bettaPtr[lenSB - 2] = element<IsLogPStrided>(logP, lblIndex, elwiseP); // logP[lbl[lenS - 1]];
#if defined(CALCULATE_ALL_IN_ONE_FRAME_LOOP)
// move to the last
gradPtr += (lenT - 1) * incG;
alphaPtr += (lenT - 1) * incA;
for (auto s = lenSB - 1; s >= 0; s--) {
auto ind = s / 2; // our real index
// we forced blanks for even indexes
auto currentInd = (s % 2 == 0) ? blankIndex : element<IsLblStrided>(lbl, ind, elwiseS);
// alpha(s)*betta(s) in log scale but still store in alpha to save memory
auto alphaBettaS = alphaPtr[s] + bettaPtr[s];
// sum (alpha(s)*betta(s) ) over real indexes
auto &currentGrad = element<isGradStrided>(gradPtr, currentInd, elwiseG); // gradPtr[currentInd];
if (currentGrad == negInf) {
currentGrad = alphaBettaS;
} else {
Type cMax = std::max(currentGrad, alphaBettaS);
currentGrad = std::log(std::exp(currentGrad - cMax) + std::exp(alphaBettaS - cMax)) + cMax;
}
}
for (int k = 0; k < lenK; k++) {
// compute the rest grad
// prob(t,k) - grad(k) / ((prob(t,k)*Z) )
// p2= grad(k) / (prob(t,k)*Z )
// in logscale . plus we have Z as -logLoss
// auto p2 = std::exp(gradPtr[k] + forwardLogLoss - logP[k]);
// gradPtr[k] = std::exp(logP[k]) - p2;
auto currentProb = element<IsLogPStrided>(logP, k, elwiseP);
auto &currentGrad = element<isGradStrided>(gradPtr, k, elwiseG);
auto p2 = std::exp(currentGrad + forwardLogLoss - currentProb);
currentGrad = std::exp(currentProb) - p2;
}
gradPtr -= incG;
alphaPtr -= incA;
#endif
auto bettaPrevPtr = bettaPtr;
bettaPtr -= incA;
logP -= incP;
// process the rest
for (auto t = lenT - 2; t >= 0; t--) {
#if defined(CALCULATE_ALL_IN_ONE_FRAME_LOOP)
auto end = lenSB - 1;
#else
auto end = std::min(2 * t + 2, lenSB - 1);
#endif
for (auto s = end; s >= 0; s--) {
auto ind = s / 2; // our real index
// we forced blanks for even indexes
auto currentInd = (s % 2 == 0) ? blankIndex : element<IsLblStrided>(lbl, ind, elwiseS); // lbl[ind];
// {t-1,s}
Type bettaS = bettaPrevPtr[s];
Type bettaS_1 = s < lenSB - 1 ? bettaPrevPtr[s + 1] : negInf;
// logP[currentInd]
auto currentProb = element<IsLogPStrided>(logP, currentInd, elwiseP);
// if blank or the same as previous
if (s < lenSB - 2 && currentInd != blankIndex && currentInd != element<IsLblStrided>(lbl, ind + 1, elwiseS)) {
Type bettaS_2 = bettaPrevPtr[s + 2];
bettaPtr[s] = log_sum_exp(bettaS, bettaS_1, bettaS_2) + currentProb;
} else {
bettaPtr[s] = log_sum_exp(bettaS, bettaS_1) + currentProb;
}
#if defined(CALCULATE_ALL_IN_ONE_FRAME_LOOP)
// alpha(s)*betta(s) in log scale but still store in alpha to save memory
auto alphaBettaS = alphaPtr[s] + bettaPtr[s];
// sum (alpha(s)*betta(s) ) over real indexes
auto &currentGrad = element<isGradStrided>(gradPtr, currentInd, elwiseG); // gradPtr[currentInd];
if (currentGrad == negInf) {
currentGrad = alphaBettaS;
} else {
Type cMax = std::max(currentGrad, alphaBettaS);
currentGrad = std::log(std::exp(currentGrad - cMax) + std::exp(alphaBettaS - cMax)) + cMax;
}
#endif
}
#if defined(CALCULATE_ALL_IN_ONE_FRAME_LOOP)
for (int k = 0; k < lenK; k++) {
// compute the rest grad
// prob(t,k) - grad(k) / ((prob(t,k)*Z) )
// p2= grad(k) / (prob(t,k)*Z )
// in logscale . plus we have Z as -logLoss
// auto p2 = std::exp(gradPtr[k] + forwardLogLoss - logP[k]);
// gradPtr[k] = std::exp(logP[k]) - p2;
auto currentProb = element<IsLogPStrided>(logP, k, elwiseP);
auto &currentGrad = element<isGradStrided>(gradPtr, k, elwiseG);
auto p2 = std::exp(currentGrad + forwardLogLoss - currentProb);
currentGrad = std::exp(currentProb) - p2;
}
alphaPtr -= incA;
gradPtr -= incG;
#endif
bettaPrevPtr = bettaPtr;
bettaPtr -= incA;
logP -= incP;
}
#if !defined(CALCULATE_ALL_IN_ONE_FRAME_LOOP)
// alpha*betta
bettaPtr = origBetta;
logP = origLogP;
for (int t = 0; t < lenT; t++) {
for (int s = 0; s < lenSB; s++) {
auto ind = s / 2; // our real index
// we forced blanks for even indexes
auto currentInd = (s % 2 == 0) ? blankIndex : element<IsLblStrided>(lbl, ind, elwiseS); // lbl[ind];
// alpha(s)*betta(s) in log scale but still store in alpha to save memory
auto alphaBettaS = alphaPtr[s] + bettaPtr[s];
// sum (alpha(s)*betta(s) ) over real indexes
auto &currentGrad = element<isGradStrided>(gradPtr, currentInd, elwiseG); // gradPtr[currentInd];
if (currentGrad == negInf) {
currentGrad = alphaBettaS;
} else {
currentGrad = log_sum_exp(currentGrad, alphaBettaS);
}
// alphaPtr[s] = alphaBettaS;
}
PRAGMA_OMP_SIMD
for (int k = 0; k < lenK; k++) {
// compute the rest grad
// prob(t,k) - grad(k) / ((prob(t,k)*Z) )
// p2= grad(k) / (prob(t,k)*Z )
// in logscale . plus we have Z as -logLoss
// auto p2 = std::exp(gradPtr[k] + forwardLogLoss - logP[k]);
// gradPtr[k] = std::exp(logP[k]) - p2;
auto currentProb = element<IsLogPStrided>(logP, k, elwiseP);
auto &currentGrad = element<isGradStrided>(gradPtr, k, elwiseG);
auto p2 = std::exp(currentGrad + forwardLogLoss - currentProb);
currentGrad = std::exp(currentProb) - p2;
}
gradPtr += incG;
bettaPtr += incA;
alphaPtr += incA;
logP += incP;
}
#endif
}
/**
* Calculates ctc loss and fills gradients
* @param logP logits matrix(lenT,lenK) pointer (log soft max input of rnn)
* @param incP stride of logits for the next time frame
* @param gradPtr gradient for output
* @param incG stride of the gradient for the next time frame
* @param lbl target label
* @param lenT frame length
* @param lenK class length
* @param lenS target label length
* @param blankIndex index of the blank label in logit class
*/
template <bool IsLogPStrided = true, bool IsLblStrided = true, bool IsGradStrided = true, typename Type,
typename IndexType>
Type unitLossAndGrad(const Type *logP, int incP, Type *gradPtr, int incG, const IndexType *lbl, int lenT, int lenK,
int lenS, int blankIndex, int elwiseP = 1, int elwiseS = 1, int elwiseG = 1) {
auto lenSB = 2 * lenS + 1;
// create temp Array for holding bettaArr [lenT,lenSB]
// create temp Array for holding alphaArr [lenT,lenSB]
int bufferC = gradPtr ? 2 : 1;
NDArray *bufferArr = NDArrayFactory::create<Type>('c', {bufferC, lenT, lenSB});
auto bufferPtr = bufferArr->bufferAsT<Type>();
auto incA = bufferArr->stridesOf()[1];
auto bettaBufferPtr = bufferPtr + bufferArr->stridesOf()[0];
Type negInf = negative_infinity<Type>();
if (gradPtr) {
if (elwiseG == 1) {
PRAGMA_OMP_SIMD
for (int i = 0; i < lenK * lenT; i++) {
gradPtr[i] = negInf;
}
} else {
auto tempPtr = gradPtr;
for (int i = 0; i < lenT; i++) {
for (int j = 0; j < lenK; j++) element<false>(tempPtr, j, elwiseG) = negInf;
tempPtr += incG;
}
}
}
// set all vals to neginf
PRAGMA_OMP_SIMD
for (int i = 0; i < bufferC * lenSB * lenT; i++) {
bufferPtr[i] = negInf;
}
// forward
Type logLoss =
forward<IsLogPStrided, IsLblStrided>(bufferPtr, incA, logP, incP, lbl, lenSB, lenT, blankIndex, elwiseP, elwiseS);
// backward and gradient if gradptr supplied
if (gradPtr)
backwardAndGrad<IsLogPStrided, IsLblStrided, IsGradStrided>(logLoss, bufferPtr, bettaBufferPtr, incA, logP, incP,
gradPtr, incG, lbl, lenS, lenT, lenK, blankIndex,
elwiseP, elwiseS, elwiseG);
delete bufferArr;
return logLoss;
}
template <typename Type, typename IndexType>
void ctc_loss_(NDArray&logits, NDArray&targetLabels, NDArray&logitsLengths,
NDArray&targetLabelLengths, NDArray &logLosses, NDArray &gradients, int blankIndex) {
// lenT - input length of T
// lenS - lenght of sequence
// lenSB - length with blanks
auto lenBatch = logits.shapeOf()[0];
auto maxLenT = logits.shapeOf()[1];
auto lenK = logits.shapeOf()[2];
auto maxLenS = targetLabels.shapeOf()[1];
// get probability buffer and targetLabels buffer
auto logP = logits.bufferAsT<Type>();
auto lblPtr = targetLabels.bufferAsT<IndexType>();
auto lenTPtr = logitsLengths.bufferAsT<IndexType>();
auto lenSPtr = targetLabelLengths.bufferAsT<IndexType>();
auto batchLbl = targetLabels.stridesOf()[0];
auto batchP = logits.stridesOf()[0];
auto incP = logits.stridesOf()[1];
auto elwiseSLen = targetLabelLengths.stridesOf()[0];
auto elwiseT = logitsLengths.stridesOf()[0];
auto elwiseS = targetLabels.stridesOf()[1];
auto elwiseP = logits.stridesOf()[2];
int elwiseLL = 0;
Type *logLossPtr = nullptr;
if (!logLosses.isEmpty()) {
elwiseLL = logLosses.stridesOf()[0];
logLossPtr = logLosses.bufferAsT<Type>();
}
// defaulting blankIndex to the last class if its incorrect or -1
if (blankIndex > maxLenS || blankIndex < 0) blankIndex = maxLenS - 1;
auto func = [logP, batchP, incP, elwiseP, lenK, lenTPtr, lenSPtr, logLossPtr, lblPtr, maxLenT, maxLenS, batchLbl,
blankIndex, elwiseT, elwiseLL, elwiseSLen, elwiseS,
&gradients](uint64_t thread_id, int64_t start, int64_t stop, int64_t increment) -> void {
Type *gradPtr = nullptr;
Type resultLoss;
int batchG, incG, elwiseG;
if (!gradients.isEmpty()) {
batchG = gradients.stridesOf()[0];
incG = gradients.stridesOf()[1];
elwiseG = gradients.stridesOf()[2];
gradPtr = gradients.bufferAsT<Type>() + start * batchG;
} else {
elwiseG = 1;
}
auto logPtr = logP + start * batchP;
auto tempLblPtr = lblPtr + start * batchLbl;
if (elwiseP == 1 && elwiseS == 1 && elwiseG == 1) {
// choose ews one
for (int batchIndex = start; batchIndex < stop; batchIndex += increment) {
auto lenT = lenTPtr[batchIndex * elwiseT];
auto lenS = lenSPtr[batchIndex * elwiseSLen];
lenT = lenT > maxLenT ? maxLenT : lenT;
lenS = lenS > maxLenS ? maxLenS : lenS;
if (lenS <= 0 || lenT <= 0) {
resultLoss = negative_infinity<Type>();
} else {
if (lenS > lenT) lenS = lenT;
resultLoss = unitLossAndGrad<false, false, false, Type, IndexType>(logPtr, incP, gradPtr, incG, tempLblPtr,
lenT, lenK, lenS, blankIndex);
}
if (gradPtr) gradPtr += batchG;
if (logLossPtr) logLossPtr[batchIndex * elwiseLL] = resultLoss;
logPtr += batchP;
tempLblPtr += batchLbl;
}
} else {
// slow strided case for all 3
for (int batchIndex = start; batchIndex < stop; batchIndex += increment) {
auto lenT = lenTPtr[batchIndex * elwiseT];
auto lenS = lenSPtr[batchIndex * elwiseSLen];
lenT = lenT > maxLenT ? maxLenT : lenT;
lenS = lenS > maxLenS ? maxLenS : lenS;
if (lenS <= 0 || lenT <= 0) {
resultLoss = negative_infinity<Type>();
} else {
if (lenS > lenT) lenS = lenT;
resultLoss = unitLossAndGrad<true, true, true, Type, IndexType>(
logPtr, incP, gradPtr, incG, tempLblPtr, lenT, lenK, lenS, blankIndex, elwiseP, elwiseS, elwiseG);
}
if (gradPtr) gradPtr += batchG;
if (logLossPtr) logLossPtr[batchIndex * elwiseLL] = resultLoss;
logPtr += batchP;
tempLblPtr += batchLbl;
}
}
};
samediff::Threads::parallel_for(func, 0, lenBatch, 1);
}
void ctcLoss(graph::Context &block, NDArray&logits, NDArray&targetLabels, NDArray&logitsLengths,
NDArray&targetLabelLengths, NDArray &logLosses, NDArray &gradients, int blankIndex) {
auto logitsDType = logits.dataType();
auto targetLabelsDType = targetLabels.dataType();
BUILD_DOUBLE_SELECTOR(logits.dataType(), targetLabels.dataType(), ctc_loss_,
(logits, targetLabels, logitsLengths, targetLabelLengths, logLosses, gradients, blankIndex),
SD_FLOAT_TYPES, SD_INDEXING_TYPES);
}
BUILD_DOUBLE_TEMPLATE( void ctc_loss_,
(NDArray&logits, NDArray&targetLabels, NDArray&logitsLengths,
NDArray&targetLabelLengths, NDArray &logLosses, NDArray &gradients, int blankIndex),
SD_FLOAT_TYPES, SD_INDEXING_TYPES);
} // namespace helpers
} // namespace ops
} // namespace sd
#endif
Reference in New Issue View Git Blame Copy Permalink