提交 ce7f361a 编写于 作者: N nhzlx

fix comments

上级 df9cbabc
......@@ -20,7 +20,7 @@ namespace paddle {
namespace inference {
namespace tensorrt {
TEST(Pool2dOpConverter, main) {
void test_pool2d(bool global_pooling) {
framework::Scope scope;
std::unordered_set<std::string> parameters;
TRTConvertValidation validator(5, parameters, scope, 1 << 15);
......@@ -28,7 +28,10 @@ TEST(Pool2dOpConverter, main) {
// The ITensor's Dims should not contain the batch size.
// So, the ITensor's Dims of input and output should be C * H * W.
validator.DeclInputVar("pool2d-X", nvinfer1::Dims3(3, 4, 4));
validator.DeclOutputVar("pool2d-Out", nvinfer1::Dims3(3, 2, 2));
if (global_pooling)
validator.DeclOutputVar("pool2d-Out", nvinfer1::Dims3(3, 1, 1));
else
validator.DeclOutputVar("pool2d-Out", nvinfer1::Dims3(3, 2, 2));
// Prepare Op description
framework::OpDesc desc;
......@@ -40,7 +43,6 @@ TEST(Pool2dOpConverter, main) {
std::vector<int> strides({2, 2});
std::vector<int> paddings({0, 0});
std::string pooling_t = "max";
bool global_pooling = false;
desc.SetAttr("pooling_type", pooling_t);
desc.SetAttr("ksize", ksize);
......@@ -55,40 +57,9 @@ TEST(Pool2dOpConverter, main) {
validator.Execute(3);
}
TEST(Pool2dOpConverter, test_global_pooling) {
framework::Scope scope;
std::unordered_set<std::string> parameters;
TRTConvertValidation validator(5, parameters, scope, 1 << 15);
// The ITensor's Dims should not contain the batch size.
// So, the ITensor's Dims of input and output should be C * H * W.
validator.DeclInputVar("pool2d-X", nvinfer1::Dims3(3, 4, 4));
validator.DeclOutputVar("pool2d-Out", nvinfer1::Dims3(3, 1, 1));
// Prepare Op description
framework::OpDesc desc;
desc.SetType("pool2d");
desc.SetInput("X", {"pool2d-X"});
desc.SetOutput("Out", {"pool2d-Out"});
std::vector<int> ksize({2, 2});
std::vector<int> strides({2, 2});
std::vector<int> paddings({0, 0});
std::string pooling_t = "max";
bool global_pooling = true;
TEST(Pool2dOpConverter, normal) { test_pool2d(false); }
desc.SetAttr("pooling_type", pooling_t);
desc.SetAttr("ksize", ksize);
desc.SetAttr("strides", strides);
desc.SetAttr("paddings", paddings);
desc.SetAttr("global_pooling", global_pooling);
LOG(INFO) << "set OP";
validator.SetOp(*desc.Proto());
LOG(INFO) << "execute";
validator.Execute(3);
}
TEST(Pool2dOpConverter, test_global_pooling) { test_pool2d(true); }
} // namespace tensorrt
} // namespace inference
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册