cpu

app.py

@@ -4,7 +4,6 @@ import argparse
 import glob
 import multiprocessing as mp
 import os
-os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
 
 # fmt: off
 import sys
@@ -40,6 +39,7 @@ def setup_cfg(args):
     add_cat_seg_config(cfg)
     cfg.merge_from_file(args.config_file)
     cfg.merge_from_list(args.opts)
+    cfg.MODEL.DEVICE = "cpu"
     cfg.freeze()
     return cfg
 

cat_seg/modeling/transformer/cat_seg_predictor.py

@@ -58,7 +58,7 @@ class CATSegPredictor(nn.Module):
         if self.test_class_texts == None:
             self.test_class_texts = self.class_texts
         device = "cuda" if torch.cuda.is_available() else "cpu"
-  
+        self.device = device
         self.tokenizer = None
         if clip_pretrained == "ViT-G" or clip_pretrained == "ViT-H":
             # for OpenCLIP models
@@ -84,12 +84,12 @@ class CATSegPredictor(nn.Module):
             prompt_templates = ['A photo of a {} in the scene',]
         else:
             raise NotImplementedError
-
-        self.text_features = self.class_embeddings(self.class_texts, prompt_templates, clip_model).permute(1, 0, 2).float()
-        self.text_features_test = self.class_embeddings(self.test_class_texts, prompt_templates, clip_model).permute(1, 0, 2).float()
         
         self.clip_model = clip_model.float()
         self.clip_preprocess = clip_preprocess
+
+        self.text_features = self.class_embeddings(self.class_texts, prompt_templates, clip_model).permute(1, 0, 2).float()
+        self.text_features_test = self.class_embeddings(self.test_class_texts, prompt_templates, clip_model).permute(1, 0, 2).float()
         
         transformer = Aggregator(
             text_guidance_dim=text_guidance_dim,
@@ -161,9 +161,9 @@ class CATSegPredictor(nn.Module):
             else:
                 texts = [template.format(classname) for template in templates]  # format with class
             if self.tokenizer is not None:
-                texts = self.tokenizer(texts).cuda()
+                texts = self.tokenizer(texts).to(self.device)
             else: 
-                texts = clip.tokenize(texts).cuda()
+                texts = clip.tokenize(texts).to(self.device)
             class_embeddings = clip_model.encode_text(texts)
             class_embeddings /= class_embeddings.norm(dim=-1, keepdim=True)
             if len(templates) != class_embeddings.shape[0]:
@@ -171,5 +171,5 @@ class CATSegPredictor(nn.Module):
                 class_embeddings /= class_embeddings.norm(dim=-1, keepdim=True)
             class_embedding = class_embeddings
             zeroshot_weights.append(class_embedding)
-        zeroshot_weights = torch.stack(zeroshot_weights, dim=1).cuda()
+        zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(self.device)
         return zeroshot_weights