Merge pull request #108 from ROCmSoftwarePlatform/transpose

Transpose

src/onnx/onnx.cpp

@@ -74,6 +74,7 @@ struct onnx_parser
         add_mem_op("Unsqueeze", &onnx_parser::parse_unsqueeze);
         add_mem_op("Slice", &onnx_parser::parse_slice);
         add_mem_op("Concat", &onnx_parser::parse_concat);
+        add_mem_op("Transpose", &onnx_parser::parse_transpose);
     }
 
     template <class F>
@@ -361,6 +362,18 @@ struct onnx_parser
         return prog.add_instruction(migraph::op::add{}, img_scaled, bias_bcast);
     }
 
+    instruction_ref
+    parse_transpose(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
+    {
+        std::vector<int64_t> perm{};
+        if(contains(attributes, "perm"))
+        {
+            auto&& perm_vals = attributes["perm"].ints();
+            perm             = std::vector<int64_t>(perm_vals.begin(), perm_vals.end());
+        }
+        return prog.add_instruction(migraph::op::transpose{perm}, args.front());
+    }
+
     void parse_from(std::istream& is)
     {
         onnx::ModelProto model;

test/onnx/onnx_test.cpp

@@ -146,6 +146,18 @@ void globalmaxpool_test()
     EXPECT(p == prog);
 }
 
+void transpose_test()
+{
+    migraph::program p;
+    auto input = p.add_parameter("0", migraph::shape{migraph::shape::float_type, {1, 2, 2, 3}});
+    std::vector<int64_t> perm{0, 3, 1, 2};
+    p.add_instruction(migraph::op::transpose{perm}, input);
+
+    auto prog = migraph::parse_onnx("transpose_test.onnx");
+
+    EXPECT(p == prog);
+}
+
 int main()
 {
     pytorch_conv_bias_test();
@@ -156,4 +168,5 @@ int main()
     imagescaler_test();
     globalavgpool_test();
     globalmaxpool_test();
+    transpose_test();
 }