新增模型嵌入白盒水印支持断点续训

watermark_generate/deals/faster_rcnn_pytorch_white_embed.py

@@ -69,7 +69,10 @@ from torch import nn
         print(f'Secret:{self.secret} secret length:{self.secret_len}')
 
         # 生成随机的投影矩阵
-        self.X_random = torch.randn((self.secret_len, w_init.shape[0])).to(self.device)
+        if os.path.exists(key_path):
+            self.X_random = np.load(key_path)
+        else:
+            self.X_random = torch.randn((self.secret_len, w_init.shape[0])).to(self.device)
         self.save_tensor(self.X_random, key_path)  # 保存投影矩阵至指定位置
 
     def get_embeder_loss(self):

watermark_generate/deals/ssd_pytorch_white_embed.py

@@ -56,7 +56,10 @@ def modify_model_project(secret_label: str, project_dir: str, public_key: str):
         print(f'Secret:{self.secret} secret length:{self.secret_len}')
 
         # 生成随机的投影矩阵
-        self.X_random = torch.randn((self.secret_len, w_init.shape[0])).to(self.device)
+        if os.path.exists(key_path):
+            self.X_random = np.load(key_path)
+        else:
+            self.X_random = torch.randn((self.secret_len, w_init.shape[0])).to(self.device)
         self.save_tensor(self.X_random, key_path)  # 保存投影矩阵至指定位置
 
     def get_embeder_loss(self):

watermark_generate/deals/yolox_pytorch_white_embed.py

@@ -166,7 +166,10 @@ class ModelEncoder:
         print(f'Secret:{self.secret} secret length:{self.secret_len}')
 
         # 生成随机的投影矩阵
-        self.X_random = torch.randn((self.secret_len, w_init.shape[0])).to(self.device)
+        if os.path.exists(key_path):
+            self.X_random = np.load(key_path)
+        else:
+            self.X_random = torch.randn((self.secret_len, w_init.shape[0])).to(self.device)
         self.save_tensor(self.X_random, key_path)  # 保存投影矩阵至指定位置
 
     def get_embeder_loss(self):