solve #1

Separius · Oct 20, 2020 · f4c2f41 · f4c2f41
1 parent c6ede67
commit f4c2f41
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/convert.py b/convert.py
@@ -10,6 +10,7 @@
 parser = argparse.ArgumentParser(description='SimCLR converter')
 parser.add_argument('tf_path', type=str, help='path of the input tensorflow file (ex: model.ckpt-250228)')
 parser.add_argument('--ema', action='store_true')
+parser.add_argument('--supervised', action='store_true')
 args = parser.parse_args()
 
 
@@ -91,6 +92,10 @@ def main():
     assert model.fc.bias.shape == b.shape
     model.fc.bias.data = b
 
+    if args.supervised:
+        save_location = f'r{depth}_{width}x_sk{1 if sk_ratio != 0 else 0}{"_ema" if use_ema_model else ""}.pth'
+        torch.save({'resnet': model.state_dict(), 'head': head.state_dict()}, save_location)
+        return
     sd = {}
     for v in contrastive_vars:
         sd[v] = ckpt_reader.get_tensor(v)