experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Merge pull request #25766 from trevor-m:tmorris_tftrt_fp16mode

Browse Source 
PiperOrigin-RevId: 234217906
This commit is contained in:
TensorFlower Gardener 2019-02-15 15:09:54 -08:00
commit e2a04a904c
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
View File

@ -3883,8 +3883,12 @@ tensorflow::Status ConvertGraphDefToEngine(
builder->setMaxWorkspaceSize(max_workspace_size_bytes); builder->setMaxWorkspaceSize(max_workspace_size_bytes);
builder->setGpuAllocator(allocator); builder->setGpuAllocator(allocator);
if (precision_mode == TrtPrecisionMode::FP16) { if (precision_mode == TrtPrecisionMode::FP16) {
builder->setHalf2Mode(true); builder->setFp16Mode(true);
} else if (precision_mode == TrtPrecisionMode::INT8) { } else if (precision_mode == TrtPrecisionMode::INT8) {
// Setting FP16 mode as well allows TRT to also consider FP16 kernels and
// use them in situations where they are faster than INT8 or where INT8 is
// not supported for a given layer.
builder->setFp16Mode(true);
builder->setInt8Mode(true); builder->setInt8Mode(true);
if (use_calibration) { if (use_calibration) {
builder->setInt8Calibrator(calibrator); builder->setInt8Calibrator(calibrator);