pytorch converter

Author: markrogersjr

nnvm/include/nnvm/top/nn.h

@@ -394,6 +394,39 @@ struct GlobalPool2DParam : public dmlc::Parameter<GlobalPool2DParam> {
   }
 };
 
+
+struct AdaptiveMaxPool2DParam : public dmlc::Parameter<AdaptiveMaxPool2DParam> {
+  TShape output_size;
+  std::string layout;
+
+  DMLC_DECLARE_PARAMETER(AdaptiveMaxPool2DParam) {
+    DMLC_DECLARE_FIELD(output_size)
+      .describe("Array of output height and width");
+    DMLC_DECLARE_FIELD(layout).set_default("NCHW")
+      .describe("Dimension ordering of data and weight. Can be 'NCHW', 'NHWC', etc."
+                "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
+                "dimensions respectively. Convolution is applied on the 'H' and"
+                "'W' dimensions.");
+  }
+};
+
+
+struct AdaptiveAvgPool2DParam : public dmlc::Parameter<AdaptiveAvgPool2DParam> {
+  TShape output_size;
+  std::string layout;
+
+  DMLC_DECLARE_PARAMETER(AdaptiveAvgPool2DParam) {
+    DMLC_DECLARE_FIELD(output_size)
+      .describe("Array of output height and width");
+    DMLC_DECLARE_FIELD(layout).set_default("NCHW")
+      .describe("Dimension ordering of data and weight. Can be 'NCHW', 'NHWC', etc."
+                "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
+                "dimensions respectively. Convolution is applied on the 'H' and"
+                "'W' dimensions.");
+  }
+};
+
+
 struct UpSamplingParam : public dmlc::Parameter<UpSamplingParam> {
   int scale;
   std::string layout;

nnvm/python/nnvm/frontend/__init__.py

@@ -7,3 +7,4 @@
 from .darknet import from_darknet
 from .tensorflow import from_tensorflow
 from .caffe2 import from_caffe2
+from .pytorch import from_pytorch

nnvm/python/nnvm/frontend/pytorch/__init__.py

@@ -0,0 +1,2 @@
+r'''PyTorch->NNVM converter'''
+from .converter import from_pytorch