Method addLSTMLayers

samples/sampleCharRNN/sampleCharRNN.cpp:576–663

}

nvinfer1::ILayer* SampleCharRNNLoop::addLSTMLayers(SampleUniquePtr<nvinfer1::INetworkDefinition>& network)
{
    nvinfer1::ILayer* dataOut{nullptr};

    nvinfer1::ITensor* data = network->addInput(mParams.bindingNames.INPUT_BLOB_NAME, nvinfer1::DataType::kFLOAT,
        nvinfer1::Dims2(mParams.seqSize, mParams.dataSize));
    ASSERT(data != nullptr);

    nvinfer1::ITensor* hiddenLayers = network->addInput(mParams.bindingNames.HIDDEN_IN_BLOB_NAME,
        nvinfer1::DataType::kFLOAT, nvinfer1::Dims2(mParams.layerCount, mParams.hiddenSize));
    ASSERT(hiddenLayers != nullptr);

    nvinfer1::ITensor* cellLayers = network->addInput(mParams.bindingNames.CELL_IN_BLOB_NAME,
        nvinfer1::DataType::kFLOAT, nvinfer1::Dims2(mParams.layerCount, mParams.hiddenSize));
    ASSERT(cellLayers != nullptr);

    nvinfer1::ITensor* sequenceSize
        = network->addInput(mParams.bindingNames.SEQ_LEN_IN_BLOB_NAME, nvinfer1::DataType::kINT32, nvinfer1::Dims{});
    ASSERT(sequenceSize != nullptr);

    // convert tensorflow weight format to trt weight format
    std::array<nvinfer1::Weights, 2> rnnw{
        SampleCharRNNBase::convertRNNWeights(mWeightMap[mParams.weightNames.RNNW_L0_NAME], mParams.dataSize),
        SampleCharRNNBase::convertRNNWeights(mWeightMap[mParams.weightNames.RNNW_L1_NAME], mParams.hiddenSize)};
    std::array<nvinfer1::Weights, 2> rnnb{
        SampleCharRNNBase::convertRNNBias(mWeightMap[mParams.weightNames.RNNB_L0_NAME]),
        SampleCharRNNBase::convertRNNBias(mWeightMap[mParams.weightNames.RNNB_L1_NAME])};

    // Store the transformed weights in the weight map so the memory can be properly released later.
    mWeightMap["rnnwL0"] = rnnw[0];
    mWeightMap["rnnwL1"] = rnnw[1];
    mWeightMap["rnnbL0"] = rnnb[0];
    mWeightMap["rnnbL1"] = rnnb[1];

    nvinfer1::ITensor* maxSequenceSize
        = network->addConstant(nvinfer1::Dims{}, Weights{DataType::kINT32, &mParams.seqSize, 1})->getOutput(0);
    ASSERT(static_cast<size_t>(mParams.layerCount) <= INDICES.size());
    LstmIO lstmNext{data, nullptr, nullptr};
    std::vector<nvinfer1::ITensor*> hiddenOutputs;
    std::vector<nvinfer1::ITensor*> cellOutputs;
    nvinfer1::Dims2 dimWL0(4 * mParams.hiddenSize, mParams.dataSize);
    nvinfer1::Dims2 dimR(4 * mParams.hiddenSize, mParams.hiddenSize);
    nvinfer1::Dims dimB{1, {4 * mParams.hiddenSize}};
    nvinfer1::Dims dim0{1, {0}};
    auto extractWeights = [](nvinfer1::Weights weights, Dims start, Dims size) -> nvinfer1::Weights {
        const char* data = static_cast<const char*>(weights.values);
        int64_t shift = samplesCommon::volume(start);
        const int sizeOfElement = samplesCommon::getElementSize(weights.type);
        int64_t count = samplesCommon::volume(size);
        ASSERT(shift + count <= weights.count);
        return nvinfer1::Weights{weights.type, data + shift * sizeOfElement, count};
    };
    for (int i = 0; i < mParams.layerCount; ++i)
    {
        nvinfer1::Dims dimW = i == 0 ? dimWL0 : dimR;
        nvinfer1::ITensor* index
            = network->addConstant(nvinfer1::Dims{}, Weights{DataType::kINT32, &INDICES[i], 1})->getOutput(0);
        nvinfer1::ITensor* hidden = network->addGather(*hiddenLayers, *index, 0)->getOutput(0);