debug:cpu/gpu device settings

.gitignore

@@ -8,3 +8,6 @@ images/
 *.pyd
 *.swp
 *.__pycache__
+
+model/
+model_classification/

env.py

@@ -1,4 +1,5 @@
 import os
+import sys
 import subprocess
 import importlib.metadata
 
@@ -20,6 +21,7 @@ def config_env():
 
     name_to_command = {'segment_anything':'git+https://github.com/facebookresearch/segment-anything.git',
                 'panopticapi':'git+https://github.com/cocodataset/panopticapi.git',
+                #TODO
                 'torch':'torch --index-url https://download.pytorch.org/whl/cu118',
                 'torchvision':'torchvision --index-url https://download.pytorch.org/whl/cu118',
                 }
@@ -28,7 +30,7 @@ def config_env():
     if env_name == 'fossil': # in case pkgs installed to unexpected env during local dev
         for package, version in packages_to_install:
             package_spec = f"{package}=={version}" if version else package
-            package_spec = name_to_command[package_spec] if package_spec in name_to_command
+            package_spec = name_to_command[package_spec] if package_spec in name_to_command else package_spec
             if not is_pkg_installed(package):
                 #TODO
                 if package=='torch' or 'torchvision' or 'tensorflow':

inference_resnet.py

@@ -74,6 +74,7 @@ def get_triplet_model(input_shape = (600, 600, 3),
     backbone = backbone_class(input_shape=input_shape, include_top=False)
     if load_weights:
         model = get_model(backbone_name,input_shape=input_shape)
+        #TODO
         model.load_weights('/users/irodri15/data/irodri15/Fossils/Models/pretrained-herbarium/Resnet50v2_NO_imagenet_None_best_1600.h5')
         trw = model.layers[0].get_weights()
         backbone.set_weights(trw)

inference_sam.py

@@ -1,9 +1,8 @@
 import torch
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-if torch.cuda.is_available():
-    device = "cuda"
-    torch.cuda.set_per_process_memory_fraction(0.3, device=device)
+if device.type == "cuda":
+    torch.cuda.set_per_process_memory_fraction(0.3, device=device.index if device.index is not None else 0)
 else:
    device = "cpu"
 print(f"Torch device: {device}")