classification-main

pytorch图片分类训练框架

代码兼容性较强，使用的是一些基本的库、基础的函数
在argparse中可以选择使用wandb，能在wandb网站中生成可视化的训练过程

1，环境

torch：https://pytorch.org/get-started/previous-versions/

>pip install timm tqdm wandb opencv-python -i https://pypi.tuna.tsinghua.edu.cn/simple
>```
### 2，数据格式
>├── 数据集路径：data_path  
>    └── image：存放所有图片  
>    └── train.txt：训练图片的绝对路径(或相对data_path下路径)和类别号，  
>&emsp; &emsp; &emsp; &emsp; (如-->image/mask/0.jpg 0 2<--表示该图片类别为0和2，空类别图片无类别号)  
>&emsp; &emsp; └── val.txt：验证图片的绝对路径(或相对data_path下路径)和类别号  
>&emsp; &emsp; └── class.txt：所有的类别名称  
### 3，run.py
>模型训练时运行该文件，argparse中有对每个参数的说明
### 4，predict_pt.py
>使用训练好的pt模型预测
### 5，export_onnx.py
>将pt模型导出为onnx模型
### 6，predict_onnx.py
>使用导出的onnx模型预测
### 7，export_trt_record
>文档中有onnx模型导出为tensort模型的详细说明
### 8，predict_trt.py
>使用导出的trt模型预测
### 9，gradio_start.py
>用gradio将程序包装成一个可视化的页面，可以在网页可视化的展示
### 10，flask_start.py
>用flask将程序包装成一个服务，并在服务器上启动
### 11，flask_request.py
>以post请求传输数据调用服务
### 12，gunicorn_config.py
>用gunicorn多进程启动flask服务：gunicorn -c gunicorn_config.py flask_start:app




### 模型注意事项：

### 目录结构
```shell
├── README.md
├── bash_output.sh
├── bash_run.sh
├── bash_watermarking.sh
├── best.onnx
├── blind_watermark #图片嵌入盲水印代码
│   ├── __init__.py
│   ├── att.py
│   ├── blind_watermark.py
│   ├── bwm_core.py
│   ├── cli_tools.py
│   ├── pool.py
│   ├── recover.py
│   ├── requirements.txt
│   └── version.py
├── block
│   ├── data_get.py #加载数据集
│   ├── loss_get.py #计算损失
│   ├── lr_get.py #动态调整学习率
│   ├── metric_get.py #获取评估指标
│   ├── model_ema.py
│   ├── model_get.py #获取模型代码
│   ├── test_model_get.py #测试获取模型代码
│   ├── train_embeder.py #模型训练嵌入白盒水印流程
│   ├── train_get.py #正常模型训练代码
│   └── val_get.py #模型验证代码
├── export_onnx.py #模型定义导出onnx格式代码
├── export_trt
├── export_trt.exe
├── export_trt_record
├── flask_request.py
├── flask_start.py
├── gradio_start.py # 用gradio将程序包装成一个可视化的页面，可以在网页可视化的展示
├── gunicorn_config.py 
├── model #模型定义，模型名称即文件名
│   ├── Alexnet.py
│   ├── GoogleNet.py
│   ├── VGG19.py
│   ├── __init__.py
│   ├── badnet.py
│   ├── layer.py
│   ├── mobilenetv2.py
│   ├── resnet.py
│   ├── test.py
│   ├── timm_model.py
│   └── yolov7_cls.py
├── predict_onnx.py #onnx格式模型文件推理
├── predict_pt.py #pt格式模型文件推理
├── predict_watermark.py #加载pt格式模型文件，并且验证白盒水印推理
├── predict_trt.py
├── prune_last.pt
├── requirement
├── run.py #模型训练脚本
├── watermarking_dataset_process.py #处理数据集脚本
└── tool
    ├── check_image.py
    ├── generate_txt.py #处理数据集，增加标签描述文件
    ├── make_flip_image.py
    ├── make_txt.py 
    ├── secret_func.py #生成验证密码标签，用于对接密码机
    ├── training_embedding.py #白盒水印编解码器，用于嵌入白盒水印和标签提取
    └── watermarking_data_process.py #对数据集进行处理，嵌入黑盒水印

训练步骤

1. 将数据集目录结构组织好
2. 使用脚本生成类型标签文本 shell 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脚本进行普通训练

   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
   

   白盒水印训练

4. 修改模型定义，添加get_encode_layers(self)函数定义，返回待嵌入水印的卷积层，类似以下代码

   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]
   

5. 调用bash_run.sh脚本进行白盒水印训练

   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
   

   黑盒水印训练

6. 调用以下命令修改数据集

   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
   

7. 黑盒水印训练不改变模型，只改变训练集，故调用bash_run.sh脚本进行普通训练即可

   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