| 493 | } |
| 494 | |
| 495 | nvinfer1::ILayer* SampleCharRNNLoop::addLSTMCell(SampleUniquePtr<nvinfer1::INetworkDefinition>& network, |
| 496 | const LstmIO& inputTensors, nvinfer1::ITensor* sequenceSize, const LstmParams& params, LstmIO& outputTensors) |
| 497 | { |
| 498 | nvinfer1::ILoop* sequenceLoop = network->addLoop(); |
| 499 | sequenceLoop->addTripLimit(*sequenceSize, nvinfer1::TripLimit::kCOUNT); |
| 500 | |
| 501 | nvinfer1::ITensor* input = sequenceLoop->addIterator(*inputTensors.data)->getOutput(0); |
| 502 | nvinfer1::IRecurrenceLayer* hidden = sequenceLoop->addRecurrence(*inputTensors.hidden); |
| 503 | nvinfer1::IRecurrenceLayer* cell = sequenceLoop->addRecurrence(*inputTensors.cell); |
| 504 | |
| 505 | nvinfer1::ITensor* mmInput = network |
| 506 | ->addMatrixMultiply(*input, nvinfer1::MatrixOperation::kVECTOR, |
| 507 | *params.inputWeights, nvinfer1::MatrixOperation::kTRANSPOSE) |
| 508 | ->getOutput(0); |
| 509 | |
| 510 | nvinfer1::ITensor* mmHidden = network |
| 511 | ->addMatrixMultiply(*hidden->getOutput(0), nvinfer1::MatrixOperation::kVECTOR, |
| 512 | *params.recurrentWeights, nvinfer1::MatrixOperation::kTRANSPOSE) |
| 513 | ->getOutput(0); |
| 514 | |
| 515 | nvinfer1::ITensor* mm |
| 516 | = network->addElementWise(*mmInput, *mmHidden, nvinfer1::ElementWiseOperation::kSUM)->getOutput(0); |
| 517 | |
| 518 | nvinfer1::ITensor* bias |
| 519 | = network->addElementWise(*params.inputBias, *params.recurrentBias, nvinfer1::ElementWiseOperation::kSUM) |
| 520 | ->getOutput(0); |
| 521 | |
| 522 | nvinfer1::ITensor* gatesICFO |
| 523 | = network->addElementWise(*mm, *bias, nvinfer1::ElementWiseOperation::kSUM)->getOutput(0); |
| 524 | |
| 525 | const auto isolateGate = [&](nvinfer1::ITensor& gates, int gateIndex) -> nvinfer1::ITensor* { |
| 526 | nvinfer1::ISliceLayer* slice = network->addSlice(gates, nvinfer1::Dims{1, {gateIndex * mParams.hiddenSize}}, |
| 527 | nvinfer1::Dims{1, {mParams.hiddenSize}}, nvinfer1::Dims{1, {1}}); |
| 528 | return addReshape(network, *slice->getOutput(0), nvinfer1::Dims{1, {mParams.hiddenSize}}); |
| 529 | }; |
| 530 | |
| 531 | nvinfer1::ITensor* i |
| 532 | = network->addActivation(*isolateGate(*gatesICFO, 0), nvinfer1::ActivationType::kSIGMOID)->getOutput(0); |
| 533 | nvinfer1::ITensor* c |
| 534 | = network->addActivation(*isolateGate(*gatesICFO, 1), nvinfer1::ActivationType::kTANH)->getOutput(0); |
| 535 | nvinfer1::ITensor* f |
| 536 | = network->addActivation(*isolateGate(*gatesICFO, 2), nvinfer1::ActivationType::kSIGMOID)->getOutput(0); |
| 537 | nvinfer1::ITensor* o |
| 538 | = network->addActivation(*isolateGate(*gatesICFO, 3), nvinfer1::ActivationType::kSIGMOID)->getOutput(0); |
| 539 | |
| 540 | nvinfer1::ITensor* C |
| 541 | = network |
| 542 | ->addElementWise(*network->addElementWise(*f, *cell->getOutput(0), nvinfer1::ElementWiseOperation::kPROD) |
| 543 | ->getOutput(0), |
| 544 | *network->addElementWise(*i, *c, nvinfer1::ElementWiseOperation::kPROD)->getOutput(0), |
| 545 | nvinfer1::ElementWiseOperation::kSUM) |
| 546 | ->getOutput(0); |
| 547 | nvinfer1::ITensor* H |
| 548 | = network |
| 549 | ->addElementWise(*o, *network->addActivation(*C, nvinfer1::ActivationType::kTANH)->getOutput(0), |
| 550 | nvinfer1::ElementWiseOperation::kPROD) |
| 551 | ->getOutput(0); |
| 552 |
nothing calls this directly
no test coverage detected