add LRN; fixed pooling params

Slayerxxx · Oct 31, 2017 · dcc0c5f · dcc0c5f
1 parent 303e512
commit dcc0c5f
Showing 1 changed file with 11 additions and 3 deletions.
diff --git a/pytorch2caffe.py b/pytorch2caffe.py
@@ -1,6 +1,6 @@
 import sys
 
-sys.path.append('/data/build_caffe/caffe_rtpose/python')
+sys.path.append('/extra/caffe/build_caffe/caffe_rtpose/python')
 import caffe
 from collections import OrderedDict
 import torch.nn as nn
@@ -26,7 +26,8 @@
               'SigmoidBackward': 'Sigmoid',
               'LeakyReLUBackward': 'ReLU',
               'NegateBackward': 'Power',
-              'MulBackward': 'Eltwise'}
+              'MulBackward': 'Eltwise',
+              'SpatialCrossMapLRNFunc': 'LRN'}
 
 layer_id = 0
 
@@ -276,7 +277,8 @@ def add_layer(func):
             pooling_param['kernel_size'] = func.kernel_size[0]
             pooling_param['stride'] = func.stride[0]
             # http://netaz.blogspot.com/2016/08/confused-about-caffes-pooling-layer.html
-            padding = 0 if func.padding[0] in {0, 1} else func.padding[0]
+            padding = func.padding[0]
+            # padding = 0 if func.padding[0] in {0, 1} else func.padding[0]
             pooling_param['pad'] = padding
             layer['pooling_param'] = pooling_param
         elif parent_type == 'AvgPool2dBackward':
@@ -301,6 +303,12 @@ def add_layer(func):
             eltwise_param = OrderedDict()
             eltwise_param['operation'] = 'SUM'
             layer['eltwise_param'] = eltwise_param
+        elif parent_type == 'SpatialCrossMapLRNFunc':
+            layer['lrn_param'] = {
+                'local_size': func.size,
+                'alpha': func.alpha,
+                'beta': func.beta,
+            }
 
         layer['top'] = parent_top  # reset layer['top'] as parent_top may change
         if parent_type != 'ViewBackward':