Match the model definition in tensorflow, so that we can use pretrained tensorflow model.

ruotianluo · ruotianluo · commit 9d4c24e83c3e · 2017-09-07T19:26:46.000-07:00
And provide the script that can convert tensorflow model to pytorch model.
diff --git a/lib/nets/resnet_v1.py b/lib/nets/resnet_v1.py
@@ -118,13 +118,13 @@ def __init__(self, block, layers, num_classes=1000):
                  bias=False)
     self.bn1 = nn.BatchNorm2d(64)
     self.relu = nn.ReLU(inplace=True)
+    # Note that, tf-faster-rcnn use padding 1 maxpool instead of ceil_mode, but it doesn't affect the output much
     self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=0, ceil_mode=True) # change
     self.layer1 = self._make_layer(block, 64, layers[0])
     self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
     self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
-    self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
-    self.avgpool = nn.AvgPool2d(7)
-    self.fc = nn.Linear(512 * block.expansion, num_classes)
+    # use stride 1 for the last conv4 layer (same as tf-faster-rcnn)
+    self.layer4 = self._make_layer(block, 512, layers[3], stride=1)
 
     for m in self.modules():
       if isinstance(m, nn.Conv2d):
diff --git a/tools/convert_from_tensorflow.py b/tools/convert_from_tensorflow.py
@@ -0,0 +1,75 @@
+import tensorflow as tf
+from tensorflow.python import pywrap_tensorflow
+from collections import OrderedDict
+import re
+import torch
+
+import argparse
+parser = argparse.ArgumentParser(description='Convert tf-faster-rcnn model to pytorch-faster-rcnn model')
+parser.add_argument('--tensorflow_model',
+                    help='the path of tensorflow_model',
+                    default=None, type=str)
+
+args = parser.parse_args()
+
+reader = pywrap_tensorflow.NewCheckpointReader(args.tensorflow_model)
+var_to_shape_map = reader.get_variable_to_shape_map()
+var_dict = {k:reader.get_tensor(k) for k in var_to_shape_map.keys()}
+
+del var_dict['Variable']
+
+for k in var_dict.keys():
+    if 'Momentum' in k:
+        del var_dict[k]
+
+for k in var_dict.keys():
+    if k.find('/') >= 0:
+        var_dict['resnet' + k[k.find('/'):]] = var_dict[k]
+        del var_dict[k]
+
+dummy_replace = OrderedDict([
+                ('moving_mean', 'running_mean'),\
+                ('moving_variance', 'running_var'),\
+                ('weights', 'weight'),\
+                ('biases', 'bias'),\
+                ('conv1/BatchNorm', 'bn1'),\
+                ('conv2/BatchNorm', 'bn2'),\
+                ('conv3/BatchNorm', 'bn3'),\
+                ('bottleneck_v1/', ''),\
+                ('block', 'layer'),\
+                ('resnet/rpn_conv/3x3', 'rpn_net'),\
+                ('resnet/rpn_cls_score', 'rpn_cls_score_net'),\
+                ('resnet/cls_score', 'cls_score_net'),\
+                ('resnet/rpn_bbox_pred', 'rpn_bbox_pred_net'),\
+                ('resnet/bbox_pred', 'bbox_pred_net'),\
+                ('shortcut/weight', 'downsample.0.weight'),\
+                ('shortcut/BatchNorm', 'downsample.1'),\
+                ('gamma', 'weight'),\
+                ('beta', 'bias'),\
+                ('/', '.')])
+
+for a, b in dummy_replace.items():
+    for k in var_dict.keys():
+        if a in k:
+            var_dict[k.replace(a,b)] = var_dict[k]
+            del var_dict[k]
+
+
+for k in var_dict.keys():
+    if 'unit_' in k:
+        m = re.search('unit_(\d+)', k)
+        var_dict[k.replace(m.group(0), str(int(m.group(1)) - 1))] = var_dict[k]
+        del var_dict[k]
+
+for k in var_dict.keys():
+    if var_dict[k].ndim == 4:
+        var_dict[k] = var_dict[k].transpose((3, 2, 0, 1)).copy(order='C')
+    if var_dict[k].ndim == 2:
+        var_dict[k] = var_dict[k].transpose((1, 0)).copy(order='C')
+    # assert x[k].shape == var_dict[k].shape, k
+
+for k in var_dict.keys():
+    var_dict[k] = torch.from_numpy(var_dict[k])
+
+
+torch.save(var_dict, args.tensorflow_model[:args.tensorflow_model.find('.ckpt')]+'.pth')
