add config

config.json

@@ -6,6 +6,7 @@
     "AutoConfig": "model_segvol_single.SegVolConfig",
     "AutoModel": "model_segvol_single.SegVolModel"
   },
+  "custom_device": "cpu",
   "model_type": "segvol",
   "patch_size": [
     4,

model_segvol_single.py

@@ -9,12 +9,14 @@ class SegVolConfig(PretrainedConfig):
     def __init__(
         self,
         test_mode=True,
+        custom_device='cpu',
         # clip_model='.',
         **kwargs,
     ):
         self.spatial_size = [32, 256, 256]
         self.patch_size = [4, 16, 16]
         self.test_mode = test_mode
+        self.custom_device = custom_device
         # self.clip_model = clip_model
         super().__init__(**kwargs)
 
@@ -36,13 +38,14 @@ class SegVolModel(PreTrainedModel):
             prompt_encoder=sam_model.prompt_encoder,
             roi_size=self.config.spatial_size,
             patch_size=self.config.patch_size,
+            custom_device=self.config.custom_device,
             # clip_model=self.config.clip_model,
             test_mode=self.config.test_mode,
             )
         
-        self.processor = SegVolProcessor(spatial_size=self.config.spatial_size)
+        self.processor = SegVolProcessor(spatial_size=self.config.spatial_size, custom_device=self.config.custom_device)
 
-        self.custom_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.custom_device = self.config.custom_device
 
     def forward_test(self, 
                 image, 
@@ -125,7 +128,7 @@ class SegVolModel(PreTrainedModel):
 
 # processor
 class SegVolProcessor():
-    def __init__(self, spatial_size) -> None:
+    def __init__(self, spatial_size, custom_device) -> None:
         self.img_loader = transforms.LoadImage()
         self.transform4test = transforms.Compose(
             [
@@ -137,7 +140,7 @@ class SegVolProcessor():
             ]
         )
         self.zoom_out_transform = transforms.Resized(keys=["image", "label"], spatial_size=spatial_size, mode='nearest-exact')
-        self.custom_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.custom_device = custom_device
         self.transform4train = transforms.Compose(
         [
             # transforms.AddChanneld(keys=["image"]),
@@ -422,15 +425,16 @@ class SegVol(nn.Module):
                 prompt_encoder,
                 roi_size,
                 patch_size,
+                custom_device,
                 # clip_model,
                 test_mode=False,
                 ):
         super().__init__()
-        self.custom_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.custom_device = custom_device
         self.image_encoder = image_encoder
         self.mask_decoder = mask_decoder
         self.prompt_encoder = prompt_encoder
-        self.text_encoder = TextEncoder()
+        self.text_encoder = TextEncoder(custom_device=custom_device)
         self.feat_shape = np.array(roi_size)/np.array(patch_size)
         self.test_mode = test_mode
         self.dice_loss = BinaryDiceLoss().to(self.custom_device)
@@ -607,9 +611,9 @@ class SegVol(nn.Module):
     #     return pseudo_labels, bboxes
     
 class TextEncoder(nn.Module):
-    def __init__(self):
+    def __init__(self, custom_device):
         super().__init__()
-        self.custom_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.custom_device = custom_device
         config = CLIPTextConfig()
         self.clip_text_model = CLIPTextModel(config)
         self.tokenizer = None