Merge pull request THU-BPM#28 from KkkyleZ/main

TS_Watermark implement

config/TS.json

@@ -0,0 +1,12 @@
+{
+  "algorithm_name": "TS",
+  "gamma": 0.5 ,
+  "delta": 2.0,
+  "ckpt_path": "./watermark/ts/model/opt/init_0.25_1.75_default.pth",
+  "hash_key": 15485863,
+  "seeding_scheme":"simple_1",
+  "prefix_length": 1,
+  "z_threshold": 4.0,
+  "store_spike_ents": true
+
+}

watermark/auto_watermark.py

@@ -14,8 +14,8 @@
 
 # =========================================================================
 # AutoWatermark.py
-# Description: This is a generic watermark class that will be instantiated 
-#              as one of the watermark classes of the library when created 
+# Description: This is a generic watermark class that will be instantiated
+#              as one of the watermark classes of the library when created
 #              with the [`AutoWatermark.load`] class method.
 # =========================================================================
 
@@ -34,7 +34,8 @@
     'EXP': 'watermark.exp.EXP',
     'EXPGumbel': 'watermark.exp_gumbel.EXPGumbel',
     'EXPEdit': 'watermark.exp_edit.EXPEdit',
-    'ITSEdit': 'watermark.its_edit.ITSEdit'
+    'ITSEdit': 'watermark.its_edit.ITSEdit',
+    'TS':'watermark.ts.TS'
 }
 
 def watermark_name_from_alg_name(name):

watermark/ts/TS_networks.py

@@ -0,0 +1,85 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+class DeltaNetwork(nn.Module):
+    def __init__(self, input_dim=2048, layers=2, init_val=2.0):
+        super(DeltaNetwork, self).__init__()
+        if layers == 2:
+            self.delta = nn.Sequential(
+                nn.Linear(input_dim, 64),
+                nn.LeakyReLU(),
+                nn.Linear(64, 1)
+            )
+        elif layers == 3:
+            self.delta = nn.Sequential(
+                nn.Linear(input_dim, 256),
+                nn.LeakyReLU(),
+                nn.Linear(256, 32),
+                nn.LeakyReLU(),
+                nn.Linear(32, 1)
+            )
+        elif layers == 5:
+            self.delta = nn.Sequential(
+                nn.Linear(input_dim, 512),
+                nn.LeakyReLU(),
+                nn.Linear(512, 128),
+                nn.LeakyReLU(),
+                nn.Linear(128, 64),
+                nn.LeakyReLU(),
+                nn.Linear(64, 16),
+                nn.LeakyReLU(),
+                nn.Linear(16, 1)
+            )
+
+        for layer in self.delta:
+            if isinstance(layer, nn.Linear):
+                torch.nn.init.kaiming_normal_(layer.weight)
+                torch.nn.init.zeros_(layer.bias)
+        self.init_val = init_val
+        nn.init.constant_(self.delta[-1].bias, init_val)  # Set bias to the calculated value
+    def forward(self, x):
+        return self.delta(x)
+
+class GammaNetwork(nn.Module):
+    def __init__(self, input_dim=2048, layers=2, init_val=0.25):
+        super(GammaNetwork, self).__init__()
+        if layers == 2:
+            self.gamma = nn.Sequential(
+                nn.Linear(input_dim, 64),
+                nn.LeakyReLU(),
+                nn.Linear(64, 1),
+                nn.Sigmoid()
+            )
+        elif layers == 3:
+            self.gamma = nn.Sequential(
+                nn.Linear(input_dim, 256),
+                nn.LeakyReLU(),
+                nn.Linear(256, 32),
+                nn.LeakyReLU(),
+                nn.Linear(32, 1),
+                nn.Sigmoid()
+            )
+        elif layers == 5:
+            self.gamma = nn.Sequential(
+                nn.Linear(input_dim, 512),
+                nn.LeakyReLU(),
+                nn.Linear(512, 128),
+                nn.LeakyReLU(),
+                nn.Linear(128, 64),
+                nn.LeakyReLU(),
+                nn.Linear(64, 16),
+                nn.LeakyReLU(),
+                nn.Linear(16, 1),
+                nn.Sigmoid()
+            )
+
+        for layer in self.gamma:
+            if isinstance(layer, nn.Linear):
+                torch.nn.init.kaiming_normal_(layer.weight)
+                torch.nn.init.zeros_(layer.bias)
+        self.init_val = init_val
+        nn.init.constant_(self.gamma[-2].bias, np.log(init_val / (1 - init_val)))  # Set bias to the calculated value
+    def forward(self, x):
+        return self.gamma(x)

watermark/ts/__init__.py

@@ -0,0 +1,20 @@
+# Copyright 2024 THU-BPM MarkLLM.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# ====================================================
+# __init__.py
+# Description: Initialization file for the kgw package
+# ====================================================
+
+from .ts import TS