添加检查点文件夹

README.md

@@ -32,4 +32,50 @@
 
 ### 分支说明
 master分支为华科大提供的原始代码
-demo分支为经过修改的demo示例
+demo分支为经过修改的demo示例
+
+### 文件说明
+```shell
+.
+├── README.md
+├── bash_process.sh
+├── bash_run.sh
+├── 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
+│   ├── train_get.py
+│   └── val_get.py
+├── export_onnx.py
+├── flask_request.py
+├── flask_start.py
+├── gunicorn_config.py
+├── last.pt
+├── model
+│   ├── layer.py
+│   ├── yolov5.py
+│   └── yolov7.py
+├── predict_pt.py
+├── run.py
+├── test.ipynb
+├── tool
+│   ├── change_dir.py  #数据集中的标签文件图片地址均为绝对路径，使用该脚本进行修改
+│   ├── check_image.py
+│   ├── generate_txt.py
+│   └── make_txt.py
+└── watermarking_data_process.py
+
+```

run.py

@@ -34,19 +34,19 @@ parser.add_argument('--wandb_image_num', default=16, type=int, help='|wandb保
 parser.add_argument('--data_path', default=r'./datasets/coco_wm', type=str, help='|数据目录|')
 parser.add_argument('--input_size', default=640, type=int, help='|输入图片大小|')
 parser.add_argument('--output_class', default=80, type=int, help='|输出类别数|')
-parser.add_argument('--weight', default='last.pt', type=str, help='|已有模型的位置，没找到模型会创建剪枝/新模型|')
+parser.add_argument('--weight', default='./checkpoint/last.pt', type=str, help='|已有模型的位置，没找到模型会创建剪枝/新模型|')
 parser.add_argument('--prune', default=False, type=bool, help='|模型剪枝后再训练(部分模型有)，需要提供prune_weight|')
 parser.add_argument('--prune_weight', default='best.pt', type=str, help='|模型剪枝的参考模型，会创建剪枝模型和训练模型|')
 parser.add_argument('--prune_ratio', default=0.5, type=float, help='|模型剪枝时的保留比例|')
 parser.add_argument('--prune_save', default='prune_best.pt', type=str, help='|保存最佳模型，每轮还会保存prune_last.pt|')
-parser.add_argument('--model', default='yolov7', type=str, help='|自定义模型选择|')
+parser.add_argument('--model', default='yolov5', type=str, help='|自定义模型选择|')
 parser.add_argument('--model_type', default='n', type=str, help='|自定义模型型号|')
 parser.add_argument('--save_path', default='best.pt', type=str, help='|保存最佳模型，除此之外每轮还会保存last.pt|')
 parser.add_argument('--loss_weight', default=((1 / 3, 0.3, 0.5, 0.2), (1 / 3, 0.4, 0.4, 0.2), (1 / 3, 0.5, 0.3, 0.2)),
                     type=tuple, help='|每个输出层(从大到小排序)的权重->[总权重、边框权重、置信度权重、分类权重]|')
 parser.add_argument('--label_smooth', default=(0.01, 0.99), type=tuple, help='|标签平滑的值|')
 parser.add_argument('--epoch', default=10, type=int, help='|训练总轮数(包含之前已训练轮数)|')
-parser.add_argument('--batch', default=2, type=int, help='|训练批量大小，分布式时为总批量|')
+parser.add_argument('--batch', default=20, type=int, help='|训练批量大小，分布式时为总批量|')
 parser.add_argument('--warmup_ratio', default=0.01, type=float, help='|预热训练步数占总步数比例，最少5步，基准为0.01|')
 parser.add_argument('--lr_start', default=0.001, type=float, help='|初始学习率，adam算法，批量小时要减小，基准为0.001|')
 parser.add_argument('--lr_end_ratio', default=0.01, type=float, help='|最终学习率=lr_end_ratio*lr_start，基准为0.01|')
@@ -67,7 +67,7 @@ parser.add_argument('--iou_threshold', default=0.5, type=float, help='|指标计
 parser.add_argument('--distributed', default=False, type=bool, help='|单机多卡分布式训练，分布式训练时batch为总batch|')
 parser.add_argument('--local_rank', default=0, type=int, help='|分布式训练使用命令后会自动传入的参数|')
 args = parser.parse_args()
-args.device_number = max(torch.cuda.device_count(), 2)  # 使用的GPU数，可能为CPU
+args.device_number = max(torch.cuda.device_count(), 1)  # 使用的GPU数，可能为CPU
 # 为CPU设置随机种子
 torch.manual_seed(999)
 # 为所有GPU设置随机种子