Method parseBinaryProto

parsers/caffe/caffeParser/caffeParser.cpp:619–684

}

IBinaryProtoBlob* CaffeParser::parseBinaryProto(const char* fileName) noexcept
{
    CHECK_NULL_RET_NULL(fileName)
    using namespace google::protobuf::io;

    std::ifstream stream(fileName, std::ios::in | std::ios::binary);
    if (!stream)
    {
        RETURN_AND_LOG_ERROR(nullptr, "Could not open file " + std::string{fileName});
    }

    IstreamInputStream rawInput(&stream);
    CodedInputStream codedInput(&rawInput);
#if GOOGLE_PROTOBUF_VERSION >= 3011000
        codedInput.SetTotalBytesLimit(INT_MAX);
#else
        // Note: This WARs the very low default size limit (64MB)
        codedInput.SetTotalBytesLimit(INT_MAX, -1);
#endif

    trtcaffe::BlobProto blob;
    bool ok = blob.ParseFromCodedStream(&codedInput);
    stream.close();

    if (!ok)
    {
        RETURN_AND_LOG_ERROR(nullptr, "parseBinaryProto: Could not parse mean file");
    }

    Dims4 dims{1, 1, 1, 1};
    if (blob.has_shape())
    {
        int size = blob.shape().dim_size(), s[4] = {1, 1, 1, 1};
        for (int i = 4 - size; i < 4; i++)
        {
            assert(blob.shape().dim(i) < INT32_MAX);
            s[i] = static_cast<int>(blob.shape().dim(i));
        }
        dims = Dims4{s[0], s[1], s[2], s[3]};
    }
    else
    {
        dims = Dims4{blob.num(), blob.channels(), blob.height(), blob.width()};
    }

    const int dataSize = parserutils::volume(dims);
    assert(dataSize > 0);

    const trtcaffe::Type blobProtoDataType = CaffeWeightFactory::getBlobProtoDataType(blob);
    const auto blobProtoData = CaffeWeightFactory::getBlobProtoData(blob, blobProtoDataType, mTmpAllocs);

    if (dataSize != (int) blobProtoData.second)
    {
        std::cout << "CaffeParser::parseBinaryProto: blob dimensions don't match data size!!" << std::endl;
        return nullptr;
    }

    const int dataSizeBytes = dataSize * CaffeWeightFactory::sizeOfCaffeType(blobProtoDataType);