修改模型嵌入白盒水印支持断点续训

watermark_generate/deals/faster_rcnn_pytorch_white_embed.py

@@ -70,7 +70,7 @@ from torch import nn
 
         # 生成随机的投影矩阵
         if os.path.exists(key_path):
-            self.X_random = np.load(key_path)
+            self.X_random = torch.tensor(np.load(key_path), dtype=torch.float).to(self.device)
         else:
             self.X_random = torch.randn((self.secret_len, w_init.shape[0])).to(self.device)
         self.save_tensor(self.X_random, key_path)  # 保存投影矩阵至指定位置

watermark_generate/deals/ssd_pytorch_white_embed.py

@@ -57,7 +57,7 @@ def modify_model_project(secret_label: str, project_dir: str, public_key: str):
 
         # 生成随机的投影矩阵
         if os.path.exists(key_path):
-            self.X_random = np.load(key_path)
+            self.X_random = torch.tensor(np.load(key_path), dtype=torch.float).to(self.device)
         else:
             self.X_random = torch.randn((self.secret_len, w_init.shape[0])).to(self.device)
         self.save_tensor(self.X_random, key_path)  # 保存投影矩阵至指定位置

watermark_generate/deals/yolox_pytorch_white_embed.py

@@ -167,7 +167,7 @@ class ModelEncoder:
 
         # 生成随机的投影矩阵
         if os.path.exists(key_path):
-            self.X_random = np.load(key_path)
+            self.X_random = torch.tensor(np.load(key_path), dtype=torch.float).to(self.device)
         else:
             self.X_random = torch.randn((self.secret_len, w_init.shape[0])).to(self.device)
         self.save_tensor(self.X_random, key_path)  # 保存投影矩阵至指定位置