add config

config.json

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

model_segvol_single.py

@@ -9,15 +9,11 @@ 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)
 
 class SegVolModel(PreTrainedModel):
@@ -38,14 +34,11 @@ 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, custom_device=self.config.custom_device)
-
-        self.custom_device = self.config.custom_device
+        self.processor = SegVolProcessor(spatial_size=self.config.spatial_size)
 
     def forward_test(self, 
                 image, 
@@ -53,7 +46,8 @@ class SegVolModel(PreTrainedModel):
                 text_prompt=None, 
                 bbox_prompt_group=None, 
                 point_prompt_group=None, 
-                use_zoom=True):
+                use_zoom=True,):
+        device = image.device
         assert image.shape[0] == 1 and zoomed_image.shape[0] == 1, 'batch size should be 1'
         assert not (text_prompt is None and bbox_prompt_group is None and point_prompt_group is None), 'Drive SegVol using at least one type of prompt'
         bbox_prompt, bbox_prompt_map, point_prompt, point_prompt_map=None, None, None, None
@@ -110,7 +104,7 @@ class SegVolModel(PreTrainedModel):
         ## inference
         with torch.no_grad():
             logits_single_cropped = sliding_window_inference(
-                    image_single_cropped.to(self.custom_device), prompt_reflection,
+                    image_single_cropped.to(device), prompt_reflection,
                     self.config.spatial_size, 1, self.model, 0.5,
                     text=text_prompt,
                     use_box=bbox_prompt is not None,
@@ -128,7 +122,7 @@ class SegVolModel(PreTrainedModel):
 
 # processor
 class SegVolProcessor():
-    def __init__(self, spatial_size, custom_device) -> None:
+    def __init__(self, spatial_size) -> None:
         self.img_loader = transforms.LoadImage()
         self.transform4test = transforms.Compose(
             [
@@ -140,7 +134,6 @@ class SegVolProcessor():
             ]
         )
         self.zoom_out_transform = transforms.Resized(keys=["image", "label"], spatial_size=spatial_size, mode='nearest-exact')
-        self.custom_device = custom_device
         self.transform4train = transforms.Compose(
         [
             # transforms.AddChanneld(keys=["image"]),
@@ -217,24 +210,24 @@ class SegVolProcessor():
         item['zoom_out_label'] = item_zoom_out['label']
         return item
 
-    def point_prompt_b(self, label_single_resize, num_positive_extra=4, num_negative_extra=0):
+    def point_prompt_b(self, label_single_resize, num_positive_extra=4, num_negative_extra=0, device='cpu'):
         point, point_label = select_points(label_single_resize, num_positive_extra=num_positive_extra, num_negative_extra=num_negative_extra)
-        points_single = (point.unsqueeze(0).float().to(self.custom_device), point_label.unsqueeze(0).float().to(self.custom_device)) 
+        points_single = (point.unsqueeze(0).float().to(device), point_label.unsqueeze(0).float().to(device)) 
         binary_points_resize = build_binary_points(point, point_label, label_single_resize.shape).unsqueeze(0).unsqueeze(0)
         return points_single, binary_points_resize
     
-    def bbox_prompt_b(self, label_single_resize):
-        box_single = generate_box(label_single_resize).unsqueeze(0).float().to(self.custom_device)
+    def bbox_prompt_b(self, label_single_resize, device='cpu'):
+        box_single = generate_box(label_single_resize).unsqueeze(0).float().to(device)
         binary_cube_resize = build_binary_cube(box_single, binary_cube_shape=label_single_resize.shape).unsqueeze(0).unsqueeze(0)
         return box_single, binary_cube_resize
 
-    def dice_score(self, preds, labels):
+    def dice_score(self, preds, labels, device='cpu'):
         assert preds.shape[0] == labels.shape[0], "predict & target batch size don't match\n" + str(preds.shape) + str(labels.shape)
         predict = preds.view(1, -1)
         target = labels.view(1, -1)
         if target.shape[1] < 1e8:
-            predict = predict.to(self.custom_device)
-            target = target.to(self.custom_device)
+            predict = predict.to(device)
+            target = target.to(device)
         predict = torch.sigmoid(predict)
         predict = torch.where(predict > 0.5, 1., 0.)
         
@@ -425,20 +418,18 @@ class SegVol(nn.Module):
                 prompt_encoder,
                 roi_size,
                 patch_size,
-                custom_device,
                 # clip_model,
                 test_mode=False,
                 ):
         super().__init__()
-        self.custom_device = custom_device
         self.image_encoder = image_encoder
         self.mask_decoder = mask_decoder
         self.prompt_encoder = prompt_encoder
-        self.text_encoder = TextEncoder(custom_device=custom_device)
+        self.text_encoder = TextEncoder()
         self.feat_shape = np.array(roi_size)/np.array(patch_size)
         self.test_mode = test_mode
-        self.dice_loss = BinaryDiceLoss().to(self.custom_device)
-        self.bce_loss = BCELoss().to(self.custom_device)
+        self.dice_loss = BinaryDiceLoss()
+        self.bce_loss = BCELoss()
         self.decoder_iter = 6
 
     def forward(self, image, text=None, boxes=None, points=None, **kwargs):
@@ -459,12 +450,13 @@ class SegVol(nn.Module):
         return sl_loss
 
     def forward_decoder(self, image_embedding, img_shape, text=None, boxes=None, points=None):
+        device = image_embedding.device
         with torch.no_grad():
             if boxes is not None:
                 if len(boxes.shape) == 2:
                     boxes = boxes[:, None, :] # (B, 1, 6)
             if text is not None:
-                text_embedding = self.text_encoder(text)  # (B, 768)
+                text_embedding = self.text_encoder(text, device)  # (B, 768)
             else:
                 text_embedding = None
         sparse_embeddings, dense_embeddings = self.prompt_encoder(
@@ -487,7 +479,8 @@ class SegVol(nn.Module):
         return logits
 
     def supervised_forward(self, image, image_embedding, img_shape, training_organs, train_labels):
-        iter_points, iter_bboxes, iter_organs = self.build_prompt_label(image.shape[0], training_organs, train_labels)
+        device = image_embedding.device
+        iter_points, iter_bboxes, iter_organs = self.build_prompt_label(image.shape[0], training_organs, train_labels, device)
         # select prompt
         prompt_options = [[None, iter_points, iter_organs], [iter_bboxes, None, iter_organs], 
                         [None, None, iter_organs], [iter_bboxes, None, None], [None, iter_points, None],
@@ -517,7 +510,7 @@ class SegVol(nn.Module):
     #         sll_loss += sll_loss_dice + sll_loss_bce
     #     return sll_loss
 
-    def build_prompt_label(self, bs, training_organs, train_labels):
+    def build_prompt_label(self, bs, training_organs, train_labels, device):
         # generate prompt & label
         iter_organs = []
         iter_bboxes = []
@@ -541,10 +534,10 @@ class SegVol(nn.Module):
             iter_points_ax.append(point)
             iter_point_labels.append(point_label)
         # batched prompt
-        iter_points_ax = torch.stack(iter_points_ax, dim=0).to(self.custom_device)
-        iter_point_labels = torch.stack(iter_point_labels, dim=0).to(self.custom_device)
+        iter_points_ax = torch.stack(iter_points_ax, dim=0).to(device)
+        iter_point_labels = torch.stack(iter_point_labels, dim=0).to(device)
         iter_points = (iter_points_ax, iter_point_labels)
-        iter_bboxes = torch.stack(iter_bboxes, dim=0).float().to(self.custom_device)
+        iter_bboxes = torch.stack(iter_bboxes, dim=0).float().to(device)
         return iter_points, iter_bboxes, iter_organs
     
     # def build_pseudo_point_prompt_label(self, input_shape, seg_labels):
@@ -611,9 +604,8 @@ class SegVol(nn.Module):
     #     return pseudo_labels, bboxes
     
 class TextEncoder(nn.Module):
-    def __init__(self, custom_device):
+    def __init__(self):
         super().__init__()
-        self.custom_device = custom_device
         config = CLIPTextConfig()
         self.clip_text_model = CLIPTextModel(config)
         self.tokenizer = None
@@ -622,20 +614,20 @@ class TextEncoder(nn.Module):
         for param in self.clip_text_model.parameters():
             param.requires_grad = False
 
-    def organ2tokens(self, organ_names):
+    def organ2tokens(self, organ_names, device):
         text_list = ['A computerized tomography of a {}.'.format(organ_name) for organ_name in organ_names]
         tokens = self.tokenizer(text_list, padding=True, return_tensors="pt")
         for key in tokens.keys():
-            tokens[key] = tokens[key].to(self.custom_device)
+            tokens[key] = tokens[key].to(device)
         return tokens
     
-    def forward(self, text):
+    def forward(self, text, device):
         if text is None:
             return None
         if type(text) is str:
             # text is supposed to be list
             text = [text]
-        tokens = self.organ2tokens(text)
+        tokens = self.organ2tokens(text, device)
         clip_outputs = self.clip_text_model(**tokens)
         text_embedding = clip_outputs.pooler_output
         text_embedding = self.dim_align(text_embedding)