MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / setLayerPrecisions

Function setLayerPrecisions

samples/common/sampleEngines.cpp:507–566  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

505}
506
507void setLayerPrecisions(INetworkDefinition& network, LayerPrecisions const& layerPrecisions)
508{
509 bool const hasGlobalPrecision{layerPrecisions.find("*") != layerPrecisions.end()};
510 auto const globalPrecision = hasGlobalPrecision ? layerPrecisions.at("*") : nvinfer1::DataType::kFLOAT;
511 bool hasLayerPrecisionSkipped{false};
512 for (int32_t layerIdx = 0; layerIdx < network.getNbLayers(); ++layerIdx)
513 {
514 auto* layer = network.getLayer(layerIdx);
515 auto const layerName = layer->getName();
516 if (layerPrecisions.find(layer->getName()) != layerPrecisions.end())
517 {
518 layer->setPrecision(layerPrecisions.at(layer->getName()));
519 }
520 else if (hasGlobalPrecision)
521 {
522 // We should not set the layer precision if its default precision is INT32 or Bool.
523 if (layer->getPrecision() == nvinfer1::DataType::kINT32
524 || layer->getPrecision() == nvinfer1::DataType::kBOOL)
525 {
526 hasLayerPrecisionSkipped = true;
527 sample::gLogVerbose << "Skipped setting precision for layer " << layerName << " because the "
528 << " default layer precision is INT32 or Bool." << std::endl;
529 continue;
530 }
531 // We should not set the constant layer precision if its weights are in INT32.
532 if (layer->getType() == nvinfer1::LayerType::kCONSTANT
533 && static_cast<IConstantLayer*>(layer)->getWeights().type == nvinfer1::DataType::kINT32)
534 {
535 hasLayerPrecisionSkipped = true;
536 sample::gLogVerbose << "Skipped setting precision for layer " << layerName << " because this "
537 << "constant layer has INT32 weights." << std::endl;
538 continue;
539 }
540 // We should not set the layer precision if the layer operates on a shape tensor.
541 if (layer->getNbInputs() >= 1 && layer->getInput(0)->isShapeTensor())
542 {
543 hasLayerPrecisionSkipped = true;
544 sample::gLogVerbose << "Skipped setting precision for layer " << layerName << " because this layer "
545 << "operates on a shape tensor." << std::endl;
546 continue;
547 }
548 if (layer->getNbInputs() >= 1 && layer->getInput(0)->getType() == nvinfer1::DataType::kINT32
549 && layer->getNbOutputs() >= 1 && layer->getOutput(0)->getType() == nvinfer1::DataType::kINT32)
550 {
551 hasLayerPrecisionSkipped = true;
552 sample::gLogVerbose << "Skipped setting precision for layer " << layerName << " because this "
553 << "layer has INT32 input and output." << std::endl;
554 continue;
555 }
556 // All heuristics passed. Set the layer precision.
557 layer->setPrecision(globalPrecision);
558 }
559 }
560
561 if (hasLayerPrecisionSkipped)
562 {
563 sample::gLogInfo << "Skipped setting precisions for some layers. Check verbose logs for more details."
564 << std::endl;

Callers 1

setupNetworkAndConfigFunction · 0.85

Calls 13

getLayerMethod · 0.80
setPrecisionMethod · 0.80
getPrecisionMethod · 0.80
isShapeTensorMethod · 0.80
findMethod · 0.45
getNbLayersMethod · 0.45
getNameMethod · 0.45
getTypeMethod · 0.45
getWeightsMethod · 0.45
getNbInputsMethod · 0.45
getInputMethod · 0.45
getNbOutputsMethod · 0.45

Tested by

no test coverage detected