#include "megbrain/opr/dnn/correlation.h"
#include "megbrain/test/autocheck.h"
#include "megbrain/test/helper.h"
#include "megbrain/test/megdnn_helper.h"
#include "megdnn/oprs.h"
#include <cmath>
#include <iomanip>
#include <random>
#include <sstream>
using namespace mgb;
namespace {
using Param = opr::CorrelationForward::Param;
void run_forward(bool is_multiply) {
RNGxorshf rng{next_rand_seed()};
using Checker = AutoOprChecker<2, 1>;
Param param;
param.format = Param::Format::NCHW;
param.is_multiply = is_multiply;
param.kernel_size = 3;
param.max_displacement = 2;
param.pad_size = 1;
param.stride1 = 2;
param.stride2 = 2;
auto make_graph = [&](const Checker::SymInpArray& inputs) -> Checker::SymOutArray {
auto o0 = opr::CorrelationForward::make(inputs[0], inputs[1], param);
return {o0};
};
auto fwd = [&](Checker::NumOutArray& dest, Checker::NumInpArray inp) {
auto opr = megdnn_naive_handle()->create_operator<megdnn::CorrelationForward>();
opr->param() = param;
auto inp_shape = inp[0]->shape();
auto num = inp_shape[0];
auto height = inp_shape[2];
auto width = inp_shape[3];
uint32_t pad_size = param.pad_size;
uint32_t kernel_size = param.kernel_size;
uint32_t stride1 = param.stride1;
uint32_t stride2 = param.stride2;
uint32_t max_displacement = param.max_displacement;
int paddedbottomheight = height + 2 * pad_size;
int paddedbottomwidth = width + 2 * pad_size;
uint32_t kernel_radius = (kernel_size - 1) / 2;
uint32_t border_size = max_displacement + kernel_radius;
uint32_t top_width =
ceil(static_cast<float>(paddedbottomwidth - border_size * 2) /
static_cast<float>(stride1));
uint32_t top_height =
ceil(static_cast<float>(paddedbottomheight - border_size * 2) /
static_cast<float>(stride1));
uint32_t neighborhood_grid_radius = max_displacement / stride2;
uint32_t neighborhood_grid_width = neighborhood_grid_radius * 2 + 1;
uint32_t top_channels = neighborhood_grid_width * neighborhood_grid_width;
megdnn::TensorShape target_shape{num, top_channels, top_height, top_width};
dest[0].dtype(dtype::Float32())
.comp_node(inp[0]->comp_node())
.resize(target_shape);
opr->exec(inp[0]->as_megdnn(), inp[1]->as_megdnn(), dest[0].as_megdnn(), {});
};
auto rand_real = [&](float lo, float hi) {
std::uniform_real_distribution<float> dist(lo, hi);
return dist(rng);
};
auto gen_inp1 = [&](HostTensorND& inp) {
auto ptr = inp.ptr<float>();
for (size_t i = 0; i < inp.shape().total_nr_elems(); ++i) {
ptr[i] = rand_real(0.06f, 0.1f);
};
};
auto gen_inp2 = [&](HostTensorND& inp) {
auto ptr = inp.ptr<float>();
for (size_t i = 0; i < inp.shape().total_nr_elems(); ++i) {
ptr[i] = rand_real(0.01f, 0.04f);
};
};
Checker::RunOptions option;
option.numdiff_eps = 1e-3;
option.numdiff_max_err = 1e-2;
Checker checker{make_graph, fwd};
checker.set_input_generator(0, gen_inp1);
checker.set_input_generator(1, gen_inp2);
checker.run({TensorShape{1, 1, 10, 10}, TensorShape{1, 1, 10, 10}}, option)
.run({TensorShape{1, 3, 50, 50}, TensorShape{1, 3, 50, 50}}, option)
.run({TensorShape{1, 1, 100, 100}, TensorShape{1, 1, 100, 100}}, option);
}
TEST(TestOprDNN, CorrelationForwardMultiply) {
REQUIRE_GPU(1);
run_forward(true);
}
TEST(TestOprDNN, CorrelationForwardSubstract) {
REQUIRE_GPU(1);
run_forward(false);
}
}