"git@developer.sourcefind.cn:gaoqiong/migraphx.git" did not exist on "6a385c2693b846e37af4fb5b706f6dc7134b2f26"
Commit 5086b1cb authored by Khalique's avatar Khalique
Browse files

add transpose parsing and tf test

parent ebfe9735
......@@ -175,6 +175,7 @@ struct tf_parser
add_mem_op("Softmax", &tf_parser::parse_softmax);
add_mem_op("Squeeze", &tf_parser::parse_squeeze, false);
add_mem_op("StridedSlice", &tf_parser::parse_stridedslice);
add_mem_op("Transpose", &tf_parser::parse_transpose, false);
}
template <class F>
......@@ -741,6 +742,15 @@ struct tf_parser
return to_nhwc(prog.add_instruction(op::squeeze{squeeze_axes}, l0));
}
instruction_ref parse_transpose(const std::string&, const attribute_map&, std::vector<instruction_ref> args)
{
auto perm = args[1]->eval().get<int32_t>().to_vector();
op::transpose op;
op.dims = std::vector<int64_t>(perm.begin(), perm.end());
return prog.add_instruction(op, args.front());
}
void parse_graph(const tensorflow::GraphDef& graph)
{
nodes = get_nodes(graph, input_nodes);
......
......@@ -389,4 +389,16 @@ TEST_CASE(tanh_test)
EXPECT(p == prog);
}
TEST_CASE(transpose_test)
{
migraphx::program p;
auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 3, 16, 16}});
migraphx::shape s0{migraphx::shape::int32_type, {4}};
p.add_literal(migraphx::literal{s0, {0,2,3,1}});
p.add_instruction(migraphx::op::transpose{{0,2,3,1}}, l0);
auto prog = optimize_tf("transpose_test.pb", false);
EXPECT(p == prog);
}
int main(int argc, const char* argv[]) { test::run(argc, argv); }
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment