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

Merge pull request #24674 from trevor-m:tmorris_tftrt_dilated_conv

Browse Source 
PiperOrigin-RevId: 227785832
This commit is contained in:
TensorFlower Gardener 2019-01-03 19:11:03 -08:00
commit 31c6f4b715
4 changed files with 483 additions and 42 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/contrib/tensorrt/BUILD
View File

@ -491,6 +491,7 @@ cuda_py_tests(
"test/binary_tensor_weight_broadcast_test.py",
"test/concatenation_test.py",
"test/const_broadcast_test.py",
"test/conv2d_test.py",
"test/identity_output_test.py",
"test/manual_test.py",
"test/memory_alignment_test.py",

97
tensorflow/contrib/tensorrt/convert/convert_nodes.cc
View File

@ -1556,6 +1556,11 @@ enum class ConvolutionType { DEFAULT, DEPTHWISE_CONV };
tensorflow::Status ConvertConv2DHelper(OpConverterParams* params, int group) {
const auto& inputs = params->inputs;
const auto& node_def = params->node_def;
if (inputs.size() != 2) {
return tensorflow::errors::InvalidArgument("Two inputs are expected for ",
node_def.op(), ", at ",
node_def.name());
}
if (inputs.at(0).is_weights()) {
return tensorflow::errors::Unimplemented(
node_def.op(), " is only implemented for tensors, not weights, at ",
@ -1567,39 +1572,61 @@ tensorflow::Status ConvertConv2DHelper(OpConverterParams* params, int group) {
node_def.name());
}
TRT_ShapedWeights weights_rsck = inputs.at(1).weights();
VLOG(2) << "weight shape: " << weights_rsck.DebugString();
if (weights_rsck.shape_.nbDims != 4) {
return tensorflow::errors::Internal(
"Conv2D expects kernel of dimension 4, at: " + node_def.name());
return tensorflow::errors::InvalidArgument(
"Conv2D expects kernel of dimension 4, at " + node_def.name());
}
TFAttrs attrs(node_def);
auto data_format = attrs.get<string>("data_format");
int c_index = (data_format == "NHWC") ? 3 : 1;
int h_index = (data_format == "NHWC") ? 1 : 2;
int w_index = (data_format == "NHWC") ? 2 : 3;
auto tf_dilations = attrs.get<std::vector<int>>("dilations");
if (tf_dilations.size() != 4) {
return tensorflow::errors::InvalidArgument(
"Convolution dilations field must specify 4 dimensions, at ",
node_def.name());
}
if (tf_dilations[0] != 1 || tf_dilations[c_index] != 1) {
return tensorflow::errors::Unimplemented(
"Dilation rate must be 1 for batch and channel dimensions, at ",
node_def.name());
}
const nvinfer1::DimsHW dilation(tf_dilations[h_index], tf_dilations[w_index]);
const auto tf_stride = attrs.get<std::vector<int>>("strides");
if (tf_stride.size() != 4) {
return tensorflow::errors::InvalidArgument(
"Convolution strides field must specify 4 dimensions, at ",
node_def.name());
}
if (tf_stride[0] != 1 || tf_stride[c_index] != 1) {
return tensorflow::errors::Unimplemented(
"Stride must be 1 for batch and channel dimensions, at ",
node_def.name());
}
const nvinfer1::DimsHW stride(tf_stride[h_index], tf_stride[w_index]);
if (params->validation_only) return tensorflow::Status::OK();
const nvinfer1::ITensor* tensor = inputs.at(0).tensor();
TFAttrs attrs(node_def);
int h_index = 2;
int w_index = 3;
auto data_format = attrs.get<string>("data_format");
if (data_format == "NHWC") {
// Transpose to NCHW (NCHW is required for IConvLayer).
const bool need_transpose = (data_format == "NHWC");
if (need_transpose) {
TF_RETURN_IF_ERROR(params->converter->TransposeTensor(
const_cast<nvinfer1::ITensor*>(tensor), {0, 3, 1, 2}, &tensor));
h_index = 1;
w_index = 2;
// TODO(jie): transpose it
}
// tensor after transpose (NCHW)
// Dimensions of transposed tensor.
const auto tensor_dim = tensor->getDimensions();
int num_groups = group;
if (num_groups == 0) num_groups = tensor_dim.d[0]; // depthwise convolution
VLOG(2) << "groups count: " << num_groups;
// For depthwise convolution, group will be 0 so set num_groups to size of
// input's channel dim. For a non-depthwise conv, num_groups will be 1.
const int num_groups = (group == 0) ? tensor_dim.d[0] : group;
if (params->converter->precision_mode() == FP16MODE) {
weights_rsck =
ConvertFP32ToFP16(params->weight_store, inputs.at(1).weights());
}
TRT_ShapedWeights weights =
params->weight_store->GetTempWeights(weights_rsck);
ReorderRSCKToKCRS(weights_rsck, &weights, num_groups);
@ -1608,35 +1635,22 @@ tensorflow::Status ConvertConv2DHelper(OpConverterParams* params, int group) {
nvinfer1::DimsHW kernel_size;
kernel_size.h() = weights.shape_.d[2];
kernel_size.w() = weights.shape_.d[3];
VLOG(2) << "RSCK: " << weights.DebugString();
VLOG(2) << "kernel size: " << kernel_size.h() << ", " << kernel_size.w();
// TODO(jie): stride. (NHWC/NCHW)
const auto tf_stride = attrs.get<std::vector<int>>("strides");
VLOG(2) << "h_INDEX" << h_index << ", w_index " << w_index;
VLOG(2) << "stride: " << tf_stride[0] << tf_stride[1] << tf_stride[2]
<< tf_stride[3];
const nvinfer1::DimsHW stride(tf_stride[h_index], tf_stride[w_index]);
// Add padding.
std::vector<std::pair<int, int>> padding;
// TODO(jie): padding.
if (attrs.get<string>("padding") == "SAME") {
// This is NCHW tensor with no batch dimension.
// 1 -> h
// 2 -> w
nvinfer1::DimsHW effective_kernel_size = kernel_size;
effective_kernel_size.h() += (kernel_size.h() - 1) * (dilation.h() - 1);
effective_kernel_size.w() += (kernel_size.w() - 1) * (dilation.w() - 1);
padding = CreateSamePadding(
stride, kernel_size,
stride, effective_kernel_size,
{static_cast<int>(tensor_dim.d[1]), static_cast<int>(tensor_dim.d[2])});
} else {
padding = {{0, 0}, {0, 0}};
}
if (padding[0].first != padding[0].second ||
padding[1].first != padding[1].second) {
// TODO(jie): handle asymmetric padding
VLOG(2) << "Padding!!!: " << padding[0].first << padding[0].second
<< padding[1].first << padding[1].second;
VLOG(2) << "TENSOR before: " << DebugString(tensor->getDimensions());
// Handle asymmetric padding.
auto pad_layer = params->converter->network()->addPadding(
*const_cast<nvinfer1::ITensor*>(tensor),
nvinfer1::DimsHW(padding[0].first, padding[1].first),
@ -1646,24 +1660,23 @@ tensorflow::Status ConvertConv2DHelper(OpConverterParams* params, int group) {
const_cast<nvinfer1::ITensor*>(tensor), pad_layer->getOutput(0));
padding = {{0, 0}, {0, 0}};
tensor = pad_layer->getOutput(0);
VLOG(2) << "TENSOR after: " << DebugString(tensor->getDimensions());
}
// Add convolution.
nvinfer1::IConvolutionLayer* layer =
params->converter->network()->addConvolution(
*const_cast<nvinfer1::ITensor*>(tensor), noutput, kernel_size,
weights.GetTrtWeights(), biases.GetTrtWeights());
TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name());
layer->setStride(stride);
layer->setPadding({padding[0].first, padding[1].first});
layer->setName(node_def.name().c_str());
layer->setNbGroups(num_groups);
layer->setDilation(dilation);
const nvinfer1::ITensor* output_tensor = layer->getOutput(0);
VLOG(2) << "TENSOR out: " << DebugString(output_tensor->getDimensions());
VLOG(2) << "data_format: " << data_format;
if (data_format == "NHWC") {
// TODO(jie): transpose it back!
// Restore transpose.
if (need_transpose) {
TF_RETURN_IF_ERROR(params->converter->TransposeTensor(
const_cast<nvinfer1::ITensor*>(output_tensor), {0, 2, 3, 1},
&output_tensor));

236
tensorflow/contrib/tensorrt/convert/convert_nodes_test.cc
View File

@ -2378,6 +2378,8 @@ TEST_F(OpConverterTest, ConvertStridedSlice) {
};
{
// Input is weights, should fail.
Reset();
NodeDef node_def = get_strided_slice_nodedef();
AddTestWeights<int32>("input", {1, 2, 3}, {1, 2, 3, 4, 5, 6});
AddTestWeights<int32>("begin", {4}, {0, 0, 0, 0});
@ -2619,6 +2621,240 @@ TEST_F(OpConverterTest, ConvertStridedSlice) {
}
}
TEST_F(OpConverterTest, ConvertConv2D) {
{
// Input list is empty, should fail.
NodeDef node_def = MakeNodeDef("my_conv2d", "Conv2D", {});
RunValidationAndConversion(
node_def, error::INVALID_ARGUMENT,
"Two inputs are expected for Conv2D, at my_conv2d");
}
// Get nodedef for Conv2D layer.
auto get_conv2d_nodedef =
[](std::vector<int> strides = {1, 1, 1, 1}, string padding = "SAME",
string data_format = "NCHW",
std::vector<int> dilations = {1, 1, 1, 1}) -> NodeDef {
Scope s = Scope::NewRootScope();
auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT);
auto filter = ops::Placeholder(s.WithOpName("weights"), DT_FLOAT);
ops::Conv2D::Attrs attrs =
ops::Conv2D::Attrs().DataFormat(data_format).Dilations(dilations);
auto conv2d = ops::Conv2D(s.WithOpName("my_conv2d"), input, filter, strides,
padding, attrs);
return conv2d.operation.node()->def();
};
{
// Input is weights, should fail.
Reset();
NodeDef node_def = get_conv2d_nodedef();
AddTestWeights<float>("input", {1, 2, 3}, {1, 2, 3, 4, 5, 6});
AddTestWeights<float>("weights", {3, 3, 1, 1}, {1, 2, 3, 4, 5, 6, 7, 8, 9});
RunValidationAndConversion(
node_def, error::UNIMPLEMENTED,
"Conv2D is only implemented for tensors, not weights, at my_conv2d");
}
{
// Filter is tensor, should fail.
Reset();
NodeDef node_def = get_conv2d_nodedef();
AddTestTensor("input", {1, 2, 3});
AddTestTensor("weights", {3, 3, 1, 1});
RunValidationAndConversion(
node_def, error::UNIMPLEMENTED,
"Kernel for Conv2D must be constant weights, at my_conv2d");
}
{
// Filter is not 4D, should fail.
Reset();
NodeDef node_def = get_conv2d_nodedef();
AddTestTensor("input", {1, 2, 3});
AddTestWeights<float>("weights", {3, 3, 1}, {1, 2, 3, 4, 5, 6, 7, 8, 9});
RunValidationAndConversion(
node_def, error::INVALID_ARGUMENT,
"Conv2D expects kernel of dimension 4, at my_conv2d");
}
{
// Dilations is not 4D, should fail.
Reset();
NodeDef node_def =
get_conv2d_nodedef({1, 1, 1, 1}, "SAME", "NCHW", {1, 1, 1});
AddTestTensor("input", {1, 2, 3});
AddTestWeights<float>("weights", {3, 3, 1, 1}, {1, 2, 3, 4, 5, 6, 7, 8, 9});
RunValidationAndConversion(
node_def, error::INVALID_ARGUMENT,
"Convolution dilations field must specify 4 dimensions, at my_conv2d");
}
{
// Dilation value is not 1 for channel, should fail.
Reset();
NodeDef node_def =
get_conv2d_nodedef({1, 1, 1, 1}, "SAME", "NCHW", {1, 2, 1, 1});
AddTestTensor("input", {1, 2, 3});
AddTestWeights<float>("weights", {3, 3, 1, 1}, {1, 2, 3, 4, 5, 6, 7, 8, 9});
RunValidationAndConversion(node_def, error::UNIMPLEMENTED,
"Dilation rate must be 1 for batch and channel "
"dimensions, at my_conv2d");
}
{
// Dilation value is not 1 for channel (NHWC), should fail.
Reset();
NodeDef node_def =
get_conv2d_nodedef({1, 1, 1, 1}, "SAME", "NHWC", {1, 1, 1, 2});
AddTestTensor("input", {2, 3, 1});
AddTestWeights<float>("weights", {3, 3, 1, 1}, {1, 2, 3, 4, 5, 6, 7, 8, 9});
RunValidationAndConversion(node_def, error::UNIMPLEMENTED,
"Dilation rate must be 1 for batch and channel "
"dimensions, at my_conv2d");
}
{
// Strides is not 4D, should fail.
Reset();
NodeDef node_def =
get_conv2d_nodedef({1, 1, 1}, "SAME", "NCHW", {1, 1, 1, 1});
AddTestTensor("input", {1, 2, 3});
AddTestWeights<float>("weights", {3, 3, 1, 1}, {1, 2, 3, 4, 5, 6, 7, 8, 9});
RunValidationAndConversion(
node_def, error::INVALID_ARGUMENT,
"Convolution strides field must specify 4 dimensions, at my_conv2d");
}
{
// Stride value is not 1 for channel, should fail.
Reset();
NodeDef node_def =
get_conv2d_nodedef({1, 2, 1, 1}, "SAME", "NCHW", {1, 1, 1, 1});
AddTestTensor("input", {1, 2, 3});
AddTestWeights<float>("weights", {3, 3, 1, 1}, {1, 2, 3, 4, 5, 6, 7, 8, 9});
RunValidationAndConversion(
node_def, error::UNIMPLEMENTED,
"Stride must be 1 for batch and channel dimensions, at my_conv2d");
}
struct TestParams {
TestParams(const std::vector<int>& input_dims,
const std::vector<float>& input,
const std::vector<int>& filter_dims,
const std::vector<float>& filter,
const std::vector<int>& strides, const string& padding,
const string& data_format, const std::vector<int>& dilations,
const std::vector<int>& expected_output_dims,
const std::vector<float>& expected_output)
: input_dims(input_dims),
input(input),
filter_dims(filter_dims),
filter(filter),
strides(strides),
padding(padding),
data_format(data_format),
dilations(dilations),
expected_output_dims(expected_output_dims),
expected_output(expected_output) {}
std::vector<int> input_dims;
std::vector<float> input;
std::vector<int> filter_dims;
std::vector<float> filter;
std::vector<int> strides;
string padding;
string data_format;
std::vector<int> dilations;
std::vector<int> expected_output_dims;
std::vector<float> expected_output;
};
// Ok.
const int kConv2DOKCases = 6;
TestParams ok_params[kConv2DOKCases] = {
// Basic
TestParams{/*input_dims=*/{1, 2, 3},
/*input=*/{0, 1, 2, 3, 3, 4},
/*filter_dims=*/{1, 2, 1, 1},
/*filter=*/{-1, 1},
/*strides=*/{1, 1, 1, 1},
/*padding=*/"VALID",
/*data_format=*/"NCHW",
/*dilations=*/{1, 1, 1, 1},
/*expected_output_dims=*/{1, 2, 2},
/*expected_output=*/{1, 1, 0, 1}},
// SAME padding (Asymmetric)
TestParams{/*input_dims=*/{1, 2, 3},
/*input=*/{0, 1, 2, 3, 3, 4},
/*filter_dims=*/{1, 2, 1, 1},
/*filter=*/{-1, 1},
/*strides=*/{1, 1, 1, 1},
/*padding=*/"SAME",
/*data_format=*/"NCHW",
/*dilations=*/{1, 1, 1, 1},
/*expected_output_dims=*/{1, 2, 3},
/*expected_output=*/{1, 1, -2, 0, 1, -4}},
// SAME padding (Symmetric)
TestParams{/*input_dims=*/{1, 2, 3},
/*input=*/{0, 1, 2, 3, 3, 4},
/*filter_dims=*/{1, 3, 1, 1},
/*filter=*/{-1, 0, 1},
/*strides=*/{1, 1, 1, 1},
/*padding=*/"SAME",
/*data_format=*/"NCHW",
/*dilations=*/{1, 1, 1, 1},
/*expected_output_dims=*/{1, 2, 3},
/*expected_output=*/{1, 2, -1, 3, 1, -3}},
// NHWC
TestParams{/*input_dims=*/{2, 3, 1},
/*input=*/{0, 1, 2, 3, 3, 4},
/*filter_dims=*/{1, 2, 1, 1},
/*filter=*/{-1, 1},
/*strides=*/{1, 1, 1, 1},
/*padding=*/"VALID",
/*data_format=*/"NHWC",
/*dilations=*/{1, 1, 1, 1},
/*expected_output_dims=*/{2, 2, 1},
/*expected_output=*/{1, 1, 0, 1}},
// Dilated
TestParams{/*input_dims=*/{1, 2, 3},
/*input=*/{0, 1, 2, 3, 3, 4},
/*filter_dims=*/{1, 2, 1, 1},
/*filter=*/{-1, 1},
/*strides=*/{1, 1, 1, 1},
/*padding=*/"VALID",
/*data_format=*/"NCHW",
/*dilations=*/{1, 1, 1, 2},
/*expected_output_dims=*/{1, 2, 1},
/*expected_output=*/{2, 1}},
// Strided
TestParams{/*input_dims=*/{1, 2, 4},
/*input=*/{0, 1, 2, 2, 3, 4, 4, 7},
/*filter_dims=*/{1, 2, 1, 1},
/*filter=*/{-1, 1},
/*strides=*/{1, 1, 1, 2},
/*padding=*/"VALID",
/*data_format=*/"NCHW",
/*dilations=*/{1, 1, 1, 1},
/*expected_output_dims=*/{1, 2, 2},
/*expected_output=*/{1, 0, 1, 3}},
};
for (int i = 0; i < kConv2DOKCases; i++) {
Reset();
NodeDef node_def =
get_conv2d_nodedef(ok_params[i].strides, ok_params[i].padding,
ok_params[i].data_format, ok_params[i].dilations);
AddTestTensor("input", ok_params[i].input_dims);
AddTestWeights<float>("weights", ok_params[i].filter_dims,
ok_params[i].filter);
RunValidationAndConversion(node_def);
TRT_TensorOrWeights output;
TF_EXPECT_OK(GetTensorOrWeights("my_conv2d", &output));
EXPECT_TRUE(output.is_tensor());
ExpectTrtDimsEqualsArray(ok_params[i].expected_output_dims,
output.tensor()->getDimensions());
std::vector<float> output_data(ok_params[i].expected_output.size());
BuildAndRun<float>({{"input", ok_params[i].input}}, "my_conv2d",
&output_data);
EXPECT_THAT(output_data, ElementsAreArray(ok_params[i].expected_output));
}
}
} // namespace convert
} // namespace tensorrt
} // namespace tensorflow

191
tensorflow/contrib/tensorrt/test/conv2d_test.py Normal file
View File

@ -0,0 +1,191 @@
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model script to test TF-TensorRT integration."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
from tensorflow.contrib.tensorrt.test import tf_trt_integration_test_base as trt_test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_nn_ops
from tensorflow.python.platform import test
def conv2d_layer(inputs,
filters,
kernel_size,
strides=(1, 1),
padding="valid",
data_format="channels_last",
dilation_rate=(1, 1),
name=None):
dtype = inputs.dtype
c_axis = -1 if data_format == "channels_last" else 1
nchan = inputs.shape[c_axis]
weights_shape = (kernel_size[0], kernel_size[1], nchan, filters)
weights = constant_op.constant(np.random.randn(*weights_shape), dtype=dtype)
padding = padding.upper()
if data_format == "channels_last":
strides = [1] + list(strides) + [1]
dilations = [1] + list(dilation_rate) + [1]
data_format = "NHWC"
else:
strides = [1, 1] + list(strides)
dilations = [1, 1] + list(dilation_rate)
data_format = "NCHW"
return gen_nn_ops.conv2d(
inputs,
weights,
strides=strides,
padding=padding,
dilations=dilations,
data_format=data_format)
def div_round_up(n, d):
return (n - 1) // d + 1
def build_graph(input_dims,
dtype,
num_filters,
data_format,
kernel_sizes,
dilation_rates,
padding="same"):
g = ops.Graph()
with g.as_default():
inp = array_ops.placeholder(
dtype=dtype, shape=[None] + input_dims[1:], name="input")
with g.device("/GPU:0"):
results = []
for kernel_size in kernel_sizes:
for dilation_rate in dilation_rates:
result = conv2d_layer(inp, num_filters, kernel_size, (1, 1), padding,
data_format, dilation_rate)
results.append(result)
output = sum(results)
output = array_ops.identity(output, name="output")
return g
class Conv2DNCHWTest(trt_test.TfTrtIntegrationTestBase):
def GetParams(self):
"""Testing conversion of Conv2D (data_format=NCHW) in TF-TRT conversion."""
np.random.seed(1234)
input_dims = [13, 3, 7, 11]
g = build_graph(
input_dims=input_dims,
dtype=dtypes.float32,
num_filters=5,
data_format="channels_first",
kernel_sizes=[(3, 3), (3, 2)],
dilation_rates=[(1, 1), (2, 3)])
return trt_test.TfTrtIntegrationTestParams(
gdef=g.as_graph_def(),
input_names=["input"],
input_dims=[input_dims],
output_names=["output"],
expected_output_dims=[(13, 5, 7, 11)])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ["TRTEngineOp_0"]
class Conv2DNHWCTest(trt_test.TfTrtIntegrationTestBase):
def GetParams(self):
"""Testing conversion of Conv2D (data_format=NCHW) in TF-TRT conversion."""
np.random.seed(1234)
input_dims = [13, 7, 11, 3]
g = build_graph(
input_dims=input_dims,
dtype=dtypes.float32,
num_filters=5,
data_format="channels_last",
kernel_sizes=[(3, 3), (3, 2)],
dilation_rates=[(1, 1), (2, 3)])
return trt_test.TfTrtIntegrationTestParams(
gdef=g.as_graph_def(),
input_names=["input"],
input_dims=[input_dims],
output_names=["output"],
expected_output_dims=[(13, 7, 11, 5)])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ["TRTEngineOp_0"]
class Conv2DStridedNCHWTest(trt_test.TfTrtIntegrationTestBase):
def GetParams(self):
"""Testing conversion of strided Conv2D (data_format=NCHW) in TF-TRT
conversion.
"""
np.random.seed(1234)
dtype = dtypes.float32
input_name = "input"
n, c, h, w = 13, 3, 7, 11
num_filters = 5
input_dims = [n, c, h, w]
output_name = "output"
g = ops.Graph()
with g.as_default():
inp = array_ops.placeholder(
dtype=dtype, shape=[None] + input_dims[1:], name=input_name)
with g.device("/GPU:0"):
output = inp
output = conv2d_layer(
output,
num_filters, (3, 2),
strides=(2, 2),
padding="same",
data_format="channels_first")
h = div_round_up(h, 2)
w = div_round_up(w, 2)
output = conv2d_layer(
output,
num_filters, (3, 3),
strides=(2, 2),
dilation_rate=(2, 3),
padding="same",
data_format="channels_first")
h = div_round_up(h, 2)
w = div_round_up(w, 2)
output = array_ops.identity(output, name=output_name)
return trt_test.TfTrtIntegrationTestParams(
gdef=g.as_graph_def(),
input_names=[input_name],
input_dims=[input_dims],
output_names=[output_name],
expected_output_dims=[(n, num_filters, h, w)])
def ExpectedEnginesToBuild(self, run_params):
"""Return the expected engines to build."""
return ["TRTEngineOp_0"]
if __name__ == "__main__":
test.main()