Function parseReshape

parsers/caffe/caffeParser/opParsers/parseReshape.cpp:24–110

namespace nvcaffeparser1
{
ILayer* parseReshape(INetworkDefinition& network, const trtcaffe::LayerParameter& msg, CaffeWeightFactory& /*weightFactory*/, BlobNameToTensor& tensors)
{
    if (!checkBlobs(msg, 1, 1))
    {
        return nullptr;
    }

    const trtcaffe::ReshapeParameter& p = msg.reshape_param();
    Dims bottomDims = tensors[msg.bottom(0)]->getDimensions();
    int axis = p.has_axis() ? p.axis() : 0;

    const ::trtcaffe::BlobShape& shape = p.shape();
    // Check that N (batch dim) is 0. TensorRT does not support reshape in batch dimension
    if (network.hasImplicitBatchDimension() && (axis == 0) && (shape.dim(0) != 0))
    {
        std::cout << "Caffe Parser: Invalid reshape param. TensorRT does not support reshape in N (batch) dimension"
                  << std::endl;
        return nullptr;
    }

    // Handle axis and dims parameters
    int axStart = std::max(0, axis - 1);
    int axEnd = p.has_num_axes()
        ? std::max(0, axis - static_cast<int>(network.hasImplicitBatchDimension()) + p.num_axes())
        : bottomDims.nbDims;
    std::vector<int> reshapeDims;

    reshapeDims.reserve(axStart);
    for (int i = 0; i < axStart; i++)
    {
        reshapeDims.push_back(bottomDims.d[i]);
    }

    for (int i = 0; i < shape.dim_size(); i++)
    {
        // skip first 0 (batch)
        if (network.hasImplicitBatchDimension() && axis == 0 && i == 0)
        {
            continue;
        }
        if (shape.dim(i) == 0)
        {
            // If there is no bottom dimension corresponding to the current axis, then the params are invalid
            assert(static_cast<int>(reshapeDims.size()) <= bottomDims.nbDims);
            reshapeDims.push_back(bottomDims.d[reshapeDims.size()]);
        }
        else
        {
            reshapeDims.push_back(shape.dim(i));
        }
    }

    for (int i = axEnd; i < bottomDims.nbDims; i++)
    {
        reshapeDims.push_back(bottomDims.d[i]);
    }

    Dims topDims{};