Merge pull request NVIDIA#693 from hXl3s/RN50/ngc-checkpoint-update

[ConvNets/PyT] Fixed distributed checkpoint loading
1024er · Sep 18, 2020 · 94518be · 94518be
2 parents a74236a + 72f40b8
commit 94518be
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/PyTorch/Classification/ConvNets/classify.py b/PyTorch/Classification/ConvNets/classify.py
@@ -63,12 +63,16 @@ def main(args):
 
     if args.weights is not None:
         weights = torch.load(args.weights)
+
+        #Temporary fix to allow NGC checkpoint loading
+        weights = {k.replace("module.", ""): v for k, v in weights.items()}
+
         model.load_state_dict(weights)
 
     model = model.cuda()
 
     if args.precision in ["AMP", "FP16"]:
-        model = network_to_half(model)
+        model = model.half()
 
 
     model.eval()

diff --git a/PyTorch/Classification/ConvNets/main.py b/PyTorch/Classification/ConvNets/main.py
@@ -363,6 +363,10 @@ def _worker_init_fn(id):
                 )
             )
             pretrained_weights = torch.load(args.pretrained_weights)
+
+            #Temporary fix to allow NGC checkpoint loading
+
+            pretrained_weights = {k.replace("module.", ""): v for k, v in pretrained_weights.items()}
         else:
             print("=> no pretrained weights found at '{}'".format(args.resume))