forked from jiuyuan/InfiniTensor
71 lines
2.5 KiB
C++
71 lines
2.5 KiB
C++
|
#include "core/graph.h"
|
||
|
#include "core/runtime.h"
|
||
|
#include "operators/pooling.h"
|
||
|
#include "test.h"
|
||
|
|
||
|
namespace infini {
|
||
|
using KDPS = vector<int>;
|
||
|
using ExpectOutput = vector<float>;
|
||
|
TEST(MaxPool, ShapeInference) {
|
||
|
Runtime cpuRuntime = CpuRuntimeObj::getInstance();
|
||
|
{
|
||
|
Graph g = make_ref<GraphObj>(cpuRuntime);
|
||
|
Tensor i = g->addTensor({1, 64, 162, 162}, DataType::UInt32);
|
||
|
const int kh = 3, kw = 3, dh = 1, dw = 1, ph = 0, pw = 0, sh = 2,
|
||
|
sw = 2;
|
||
|
auto op =
|
||
|
g->addOp<MaxPoolObj>(i, nullptr, kh, kw, dh, dw, ph, pw, sh, sw);
|
||
|
EXPECT_EQ(op->getOutput()->getDims(), (Shape{1, 64, 80, 80}));
|
||
|
}
|
||
|
|
||
|
{ // dilation & stride
|
||
|
Graph g = make_ref<GraphObj>(cpuRuntime);
|
||
|
Tensor i = g->addTensor({1, 64, 162, 162}, DataType::UInt32);
|
||
|
auto op = g->addOp<MaxPoolObj>(i, nullptr, 4, 3, 1, 1, 2, 1, 1, 2);
|
||
|
EXPECT_EQ(op->getOutput()->getDims(), (Shape{1, 64, 163, 81}));
|
||
|
}
|
||
|
}
|
||
|
|
||
|
TEST(MaxPool, NaiveCPU) {
|
||
|
Runtime cpuRuntime = CpuRuntimeObj::getInstance();
|
||
|
Graph g = make_ref<GraphObj>(cpuRuntime);
|
||
|
Tensor i = g->addTensor({1, 2, 5, 5}, DataType::UInt32);
|
||
|
auto op = g->addOp<MaxPoolObj>(i, nullptr, 3, 3, 1, 1, 1, 1, 2, 2);
|
||
|
|
||
|
g->dataMalloc();
|
||
|
i->setData(IncrementalGenerator());
|
||
|
cpuRuntime->run(g, true, true);
|
||
|
double perfTime = cpuRuntime->getPerfTime(g);
|
||
|
// The example matmul takes 0.0036ms with one core
|
||
|
EXPECT_GT(perfTime, 0);
|
||
|
EXPECT_LT(perfTime, 5);
|
||
|
// check answer
|
||
|
vector<uint32_t> ans = {6, 8, 9, 16, 18, 19, 21, 23, 24,
|
||
|
31, 33, 34, 41, 43, 44, 46, 48, 49};
|
||
|
EXPECT_TRUE(op->getOutput()->equalData(ans));
|
||
|
}
|
||
|
|
||
|
TEST(AvgPool, NaiveCPU) {
|
||
|
Runtime cpuRuntime = CpuRuntimeObj::getInstance();
|
||
|
Graph g = make_ref<GraphObj>(cpuRuntime);
|
||
|
Tensor i = g->addTensor({1, 2, 5, 5}, DataType::Float32);
|
||
|
auto op = g->addOp<AvgPoolObj>(i, nullptr, 3, 3, 1, 1, 1, 1, 2, 2);
|
||
|
|
||
|
g->dataMalloc();
|
||
|
i->setData(IncrementalGenerator());
|
||
|
cpuRuntime->run(g, true, true);
|
||
|
|
||
|
// check answer
|
||
|
vector<float> ans = {
|
||
|
1.33333337, 3.0000, 2.66666675, 7.0000, 12.0000, 9.0000,
|
||
|
8.0000, 13.0000, 9.33333302, 12.444447, 19.666666, 13.7777777,
|
||
|
23.666666, 37.0000, 25.666666, 19.1111107, 29.666666, 20.4444447};
|
||
|
EXPECT_TRUE(op->getOutput()->equalData(ans));
|
||
|
|
||
|
double perfTime = cpuRuntime->getPerfTime(g);
|
||
|
// The example matmul takes 0.0036ms with one core
|
||
|
EXPECT_GT(perfTime, 0);
|
||
|
EXPECT_LT(perfTime, 5);
|
||
|
}
|
||
|
|
||
|
} // namespace infini
|