修改README

README.md

@@ -93,6 +93,7 @@
 ├── prune_last.pt
 ├── requirement
 ├── run.py #模型训练脚本
+├── watermarking_dataset_process.py #处理数据集脚本
 └── tool
     ├── check_image.py
     ├── generate_txt.py #处理数据集，增加标签描述文件
@@ -110,4 +111,33 @@
 python ./tool/generate_txt.py --txt_path './dataset/CIFAR-10_ori'  --specific_data 'train_cifar10_JPG' --txt_name 'train'
 python ./tool/generate_txt.py --txt_path './dataset/CIFAR-10_ori'  --specific_data 'test_cifar10_JPG'  --txt_name 'test'
 ```
-3. 调用bash_run.sh脚本进行训练
+3. 调用bash_run.sh脚本进行普通训练
+```shell
+python ./run.py --model Alexnet --input_size 32 --save_path ./checkpoints/Alexnet/clean/best.pt --save_path_last ./checkpoints/Alexnet/clean/last.pt --epoch 10 --data_path ./dataset --dataset_name CIFAR-10
+```
+### 白盒水印训练
+1. 修改模型定义，添加get_encode_layers(self)函数定义，返回待嵌入水印的卷积层，类似以下代码
+```python
+    def get_encode_layers(self):
+        """
+        获取用于白盒模型水印加密层，每个模型根据复杂度选择合适的卷积层
+        """
+        conv_list = []
+        for module in self.modules():
+            if isinstance(module, nn.Conv2d):
+                conv_list.append(module)
+        return conv_list[0:2]
+```
+2. 调用bash_run.sh脚本进行白盒水印训练
+```shell
+python ./run.py --model Alexnet --input_size 32 --save_path ./checkpoints/Alexnet/white_box_embed/best.pt --save_path_last ./checkpoints/Alexnet/white_box_embed/last.pt --epoch 10 --data_path ./dataset --dataset_name CIFAR-10 --white_box_embed True
+```
+### 黑盒水印训练
+1. 调用以下命令修改数据集
+```shell
+python ./watermarking_dataset_process.py --key_path ./dataset/watermarking2/key_hex.txt --dataset_train_txt_path ./dataset/CIFAR-10_wm/train.txt --dataset_test_txt_path ./dataset/CIFAR-10_wm/test.txt --dataset_name CIFAR-10 --key_size 256 --output_class 10
+```
+2. 黑盒水印训练不改变模型，只改变训练集，故调用bash_run.sh脚本进行普通训练即可
+```shell
+python ./run.py --model Alexnet --input_size 32 --save_path ./checkpoints/Alexnet/clean/best.pt --save_path_last ./checkpoints/Alexnet/clean/last.pt --epoch 10 --data_path ./dataset --dataset_name CIFAR-10
+```