修改白盒水印嵌入方法实现

watermark_codec/model_decoder.py

@@ -15,7 +15,7 @@ from watermark_codec.tool.tensor_deal import load_tensor, flatten_parameters, ge
 
 
 class ModelDecoder:
-    def __init__(self, layers: List[nn.Conv2d], key_path: str = None, device='cuda'):
+    def __init__(self, layers: List[nn.Conv2d], key_path: str = None):
         # 判断传入的层是否全部为卷积层
         for layer in layers:
             if not isinstance(layer, nn.Conv2d):

watermark_codec/model_encoder.py

@@ -15,8 +15,7 @@ from watermark_codec.tool.tensor_deal import flatten_parameters, save_tensor, ge
 
 
 class ModelEncoder:
-    def __init__(self, layers: List[nn.Conv2d], secret: str, key_path: str = None, device='cuda'):
-        self.device = device
+    def __init__(self, layers: List[nn.Conv2d], secret: str, key_path: str = None):
         self.layers = layers
 
         # 处理待嵌入的卷积层
@@ -25,7 +24,6 @@ class ModelEncoder:
                 raise TypeError('传入参数不是卷积层')
         weights = [x.weight for x in layers]
         w = flatten_parameters(weights)
-        w_init = w.clone().detach()
         print('Size of embedding parameters:', w.shape)
 
         # 对密钥进行处理
@@ -34,7 +32,7 @@ class ModelEncoder:
         print(f'Secret:{self.secret} secret length:{self.secret_len}')
 
         # 生成随机的投影矩阵
-        self.X_random = ms.ops.randn((self.secret_len, w_init.shape[0]))
+        self.X_random = ms.ops.randn((self.secret_len, w.shape[0]))
         save_tensor(self.X_random, key_path)  # 保存投影矩阵至指定位置
 
     def get_loss(self, loss, alpha=1):

watermark_codec/tool/tensor_deal.py

@@ -41,8 +41,7 @@ def flatten_parameters(weights: List[Tensor]) -> Tensor:
         mean_x = mnp.mean(x, axis=3).reshape(-1)
         mean_list.append(mean_x)
 
-    concat = ms.ops.Concat(1)
-    return concat(mean_list)
+    return ms.ops.concat(mean_list)
 
 
 def get_prob(x_random, w) -> Tensor:
@@ -66,6 +65,7 @@ def loss_fun(x, y) -> Tensor:
     return ms.ops.binary_cross_entropy(x, y)
 
 if __name__ == '__main__':
+    ms.set_context(device_target="GPU")
     key_path = './secret.ckpt'
     # 生成随机矩阵
     X_random = ms.ops.randn((2, 3))