fix bugs

tools/publish_model.py

@@ -1,9 +1,9 @@
 import argparse
 import subprocess
+from collections import OrderedDict
 
 import torch
 
-
 def parse_args():
     parser = argparse.ArgumentParser(
         description='Process a checkpoint to be published')
@@ -18,6 +18,13 @@ def process_checkpoint(in_file, out_file):
     # remove optimizer for smaller file size
     if 'optimizer' in checkpoint:
         del checkpoint['optimizer']
+    if 'state_dict' in checkpoint:
+        in_state_dict = checkpoint.pop('state_dict')
+        out_state_dict = OrderedDict()
+        for key, val in in_state_dict.items():
+            key = key.replace('backbone.','')
+            out_state_dict[key] = val
+        checkpoint['state_dict'] = out_state_dict
     # if it is necessary to remove some sensitive data in checkpoint['meta'],
     # add the code here.
     torch.save(checkpoint, out_file)