Update model/openllama.py

model/openllama.py

@@ -172,16 +172,16 @@ class OpenLLAMAPEFTModel(nn.Module):
         print (f'Initializing visual encoder from {imagebind_ckpt_path} ...')
 
         self.visual_encoder, self.visual_hidden_size = imagebind_model.imagebind_huge(args)
-        self.visual_encoder.to(torch.bfloat16).to(self.device)
+        self.visual_encoder.to(torch.float16).to(self.device)
         imagebind_ckpt = torch.load(imagebind_ckpt_path, map_location=torch.device('cpu'))
         self.visual_encoder.load_state_dict(imagebind_ckpt, strict=True)
         
 
         self.iter = 0
 
-        self.image_decoder = LinearLayer(1280, 1024, 4).to(torch.bfloat16).to(self.device)
+        self.image_decoder = LinearLayer(1280, 1024, 4).to(torch.float16).to(self.device)
 
-        self.prompt_learner = PromptLearner(1, 4096).to(torch.bfloat16).to(self.device)
+        self.prompt_learner = PromptLearner(1, 4096).to(torch.float16).to(self.device)
 
         self.loss_focal = FocalLoss()
         self.loss_dice = BinaryDiceLoss()
@@ -215,7 +215,7 @@ class OpenLLAMAPEFTModel(nn.Module):
         # # self.llama_model = load_checkpoint_and_dispatch(self.llama_model, vicuna_ckpt_path, device_map=device_map, offload_folder="offload", offload_state_dict = True)
         # # self.llama_model.to(torch.float16)
         # # try:
-        self.llama_model = AutoModelForCausalLM.from_pretrained(vicuna_ckpt_path, torch_dtype=torch.bfloat16, device_map='auto', load_in_8bit=True, offload_folder="offload1")
+        self.llama_model = AutoModelForCausalLM.from_pretrained(vicuna_ckpt_path, torch_dtype=torch.float16, device_map='auto', load_in_8bit=True, offload_folder="offload1")
         # # except:
         #     pass
         # finally:
@@ -225,7 +225,7 @@ class OpenLLAMAPEFTModel(nn.Module):
         self.llama_model.load_state_dict(delta_ckpt, strict=False)
         self.llama_model.print_trainable_parameters()
 
-        self.llama_tokenizer = LlamaTokenizer.from_pretrained(vicuna_ckpt_path, use_fast=False, torch_dtype=torch.bfloat16, device_map='auto', offload_folder="offload2")
+        self.llama_tokenizer = LlamaTokenizer.from_pretrained(vicuna_ckpt_path, use_fast=False, torch_dtype=torch.float16, device_map='auto', offload_folder="offload2")
         self.llama_tokenizer.pad_token = self.llama_tokenizer.eos_token
         self.llama_tokenizer.padding_side = "right"
         print ('Language decoder initialized.')
@@ -634,10 +634,10 @@ class OpenLLAMAPEFTModel(nn.Module):
                 anomaly_map = (100.0 * patch_tokens[layer] @ feats_text_tensor.transpose(-2,-1))
                 B, L, C = anomaly_map.shape
                 H = int(np.sqrt(L))
-                anomaly_map = anomaly_map.to(torch.float16)
+                # anomaly_map = anomaly_map.to(torch.float16)
                 anomaly_map = F.interpolate(anomaly_map.permute(0, 2, 1).view(B, 2, H, H),
                                             size=224, mode='bilinear', align_corners=True)
-                anomaly_map = anomaly_map.to(torch.bfloat16)
+                # anomaly_map = anomaly_map.to(torch.bfloat16)
                 anomaly_map = torch.softmax(anomaly_map, dim=1)
                 anomaly_maps.append(anomaly_map[:,1,:,:])
 
@@ -661,9 +661,9 @@ class OpenLLAMAPEFTModel(nn.Module):
                     sims.append(sim_max)
 
                 sim = torch.mean(torch.stack(sims,dim=0), dim=0).reshape(1,1,16,16)
-                anomaly_map = anomaly_map.to(torch.float16)
+                # anomaly_map = anomaly_map.to(torch.float16)
                 sim = F.interpolate(sim,size=224, mode='bilinear', align_corners=True)
-                anomaly_map = anomaly_map.to(torch.bfloat16)
+                # anomaly_map = anomaly_map.to(torch.bfloat16)
                 anomaly_map_ret = 1 - sim # (anomaly_map_ret + 1 - sim) / 2