Added ability to finetune models

export_model.py

@@ -12,10 +12,10 @@ output=32
 # load PyTorch model from .pth file
 
 device = torch.device("cpu")
-if torch.cuda.is_available():
-    device = torch.device("cuda")
+# if torch.cuda.is_available():
+#     device = torch.device("cuda")
 
-CHECKPOINT_PATH = "out-checkpoints/augment_rotate_75_x8/checkpoint_embed_1105.pth"
+CHECKPOINT_PATH = "checkpoints/checkpoint_embed_1105.pth"
 checkpoint = torch.load(CHECKPOINT_PATH, map_location=device)
 
 model = SPOTER_EMBEDDINGS(
@@ -30,7 +30,10 @@ model.eval()
 model_export = "onnx"
 if model_export == "coreml":
     dummy_input = torch.randn(1, 10, 54, 2)
+    # set device for dummy input
+    dummy_input = dummy_input.to(device)
     traced_model = torch.jit.trace(model, dummy_input)
+
     out = traced_model(dummy_input)
     import coremltools as ct
 
@@ -41,10 +44,12 @@ if model_export == "coreml":
     )
 
     # Save Core ML model
-    coreml_model.save("models/" + model_name + ".mlmodel")
+    coreml_model.save("out-models/" + model_name + ".mlmodel")
 else:
     # create dummy input tensor
     dummy_input = torch.randn(1, 10, 54, 2)
+    # set device for dummy input
+    dummy_input = dummy_input.to(device)
 
     # export model to ONNX format
     output_file = 'models/' + model_name + '.onnx'
@@ -52,7 +57,7 @@ else:
 
     torch.onnx.export(model,                                # model being run
                     dummy_input,                          # model input (or a tuple for multiple inputs)
-                    'output-models/' + model_name + '.onnx',     # where to save the model (can be a file or file-like object)
+                    'out-models/' + model_name + '.onnx',     # where to save the model (can be a file or file-like object)
                     export_params=True,                   # store the trained parameter weights inside the model file
                     opset_version=9,                      # the ONNX version to export the model to
                     do_constant_folding=True,             # whether to execute constant folding for optimization