bn without affine

Slayerxxx · Aug 29, 2017 · 82741b3 · 82741b3
1 parent 8c7e86e
commit 82741b3
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 20 deletions.
diff --git a/prototxt.py b/prototxt.py
@@ -149,7 +149,7 @@ def format_value(value):
         # if str.isnumeric():
         if is_number(value):
             return value
-        elif value == 'true' or value == 'false' or value == 'MAX' or value == 'SUM' or value == 'AVE':
+        elif value in {'true', 'false', 'MAX', 'SUM', 'AVE', 'PROD'}:
             return value
         else:
             return '\"%s\"' % value

diff --git a/pytorch2caffe.py b/pytorch2caffe.py
@@ -37,6 +37,8 @@ def pytorch2caffe(input_var, output_var, protofile, caffemodel):
     print_prototxt(net_info)
     save_prototxt(net_info, protofile)
 
+    if caffemodel is None:
+        return
     net = caffe.Net(protofile, caffe.TEST)
     params = net.params
 
@@ -73,25 +75,25 @@ def convert_layer(func):
             elif parent_type == 'BatchNormBackward':
                 running_mean = func.running_mean
                 running_var = func.running_var
-                # print('%s running_mean' % parent_name, running_mean)
-                # exit(0)
-                scale_weights = func.next_functions[1][0].variable.data
-                scale_biases = func.next_functions[2][0].variable.data
                 bn_name = parent_name + "_bn"
-                scale_name = parent_name + "_scale"
                 save_bn2caffe(running_mean, running_var, params[bn_name])
-                save_scale2caffe(scale_weights, scale_biases, params[scale_name])
+
+                affine = func.next_functions[1][0] is not None
+                if affine:
+                    scale_weights = func.next_functions[1][0].variable.data
+                    scale_biases = func.next_functions[2][0].variable.data
+                    scale_name = parent_name + "_scale"
+                    save_scale2caffe(scale_weights, scale_biases, params[scale_name])
             elif parent_type == 'AddmmBackward':
                 biases = func.next_functions[0][0].variable.data
                 weights = func.next_functions[2][0].next_functions[0][0].variable.data
                 save_fc2caffe(weights, biases, params[parent_name])
             elif parent_type == 'UpsamplingNearest2d':
                 print('UpsamplingNearest2d')
 
-    if caffemodel is not None:
-        convert_layer(output_var.grad_fn)
-        print('save caffemodel to %s' % caffemodel)
-        net.save(caffemodel)
+    convert_layer(output_var.grad_fn)
+    print('save caffemodel to %s' % caffemodel)
+    net.save(caffemodel)
 
 
 def save_conv2caffe(weights, biases, conv_param):
@@ -236,14 +238,20 @@ def add_layer(func):
             batch_norm_param['use_global_stats'] = 'true'
             bn_layer['batch_norm_param'] = batch_norm_param
 
-            scale_layer = OrderedDict()
-            scale_layer['name'] = parent_name + "_scale"
-            scale_layer['type'] = 'Scale'
-            scale_layer['bottom'] = parent_top
-            scale_layer['top'] = parent_top
-            scale_param = OrderedDict()
-            scale_param['bias_term'] = 'true'
-            scale_layer['scale_param'] = scale_param
+            affine = func.next_functions[1][0] is not None
+            # func.next_functions[1][0].variable.data
+            if affine:
+                scale_layer = OrderedDict()
+                scale_layer['name'] = parent_name + "_scale"
+                scale_layer['type'] = 'Scale'
+                scale_layer['bottom'] = parent_top
+                scale_layer['top'] = parent_top
+                scale_param = OrderedDict()
+                scale_param['bias_term'] = 'true'
+                scale_layer['scale_param'] = scale_param
+            else:
+                scale_layer = None
+
         elif parent_type == 'ThresholdBackward':
             parent_top = parent_bottoms[0]
         elif parent_type == 'MaxPool2dBackward':
@@ -282,7 +290,8 @@ def add_layer(func):
         if parent_type != 'ViewBackward':
             if parent_type == "BatchNormBackward":
                 layers.append(bn_layer)
-                layers.append(scale_layer)
+                if scale_layer is not None:
+                    layers.append(scale_layer)
             else:
                 layers.append(layer)
                 # layer_id = layer_id + 1