Skip to content

GitLab

Commit 1707011e authored by chenych's avatar chenych
Browse files

update codes

parent 201eaf53
Hide whitespace changes
Inline Side-by-side
Showing with 27 additions and 45 deletions
README.md
View file @ 1707011e
......@@ -4,7 +4,7 @@
## 模型结构
使用传统的CNN主干来学习输入图像的2D表示。该模型对其进行平坦化，并在将其传递到转换器编码器之前用位置编码对其进行补充。然后，转换器解码器将少量固定数量的学习位置嵌入作为输入，称之为对象查询，并额外处理编码器输出。将解码器的每个输出嵌入传递到共享前馈网络（FFN），该网络预测检测（类和边界框）或“无对象”类。
对于输入图像，DETR使用传统的CNN backbone来学习2D特征, 并在将其传递到transformer encoder之前用positional encoding对其进行补充和平滑。然后，transformer decoder将少量固定数量的学习到的positional embeddings(对象查询)作为输入，额外关注encoder输出。将解码器的每个输出embedding传递到预测检测（类和边界框）或“无对象”类的共享前馈网络（FFN）。
<div align=center>
<img src="./doc/models.png"/>
......@@ -104,14 +104,12 @@ bash train.sh
验证前需提前准备好预训练模型，并将coco_path设置为当前环境中推理数据的对应地址，数据应为COCO数据格式
如果想要查看预测效果（预测结果输出到图片上），请执行
如果想要查看预测效果（预测结果输出到图片上），请执行(其余参数与训练模型参数需一致，详情请根据训练参数配置):
```
python test.py --pre_trained_model <checkpoint path> --coco_path <coco path>
```
其余对应参数与训练模型参数需一致，详情请参考代码里面的参数配置
#### 单卡推理
```
......@@ -122,17 +120,17 @@ bash val.sh
<div align=center>
<img src="./doc/xxx.png"/>
<img src="./doc/000000001503.jpg"/>
</div>
### 精度
在COCO2017的测试集上进行单卡测试，结果如下表所示(待补充)
在COCO2017的val数据集上进行单卡测试，结果如下表所示:
| xxx | xxx | xxx | xxx | xxx |
| name | backbone | schedule | inf_time | box AP |
| :------: | :------: | :------: | :------: |:------: |
| xxx | xxx | xxx | xxx | xxx |
| xxx | xx | xxx | xxx | xxx |
| DETR(our) | R50 | 500 | xxx | xxx |
| [DETR](https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth) | R50 | 500 | 0.036 | 42.0 |
## 应用场景
......
docker/Dockerfile
View file @ 1707011e
FROM image.sourcefind.cn:5000/dcu/admin/base/pytorch:1.13.1-centos7.6-dtk-23.04.1-py38-latest
RUN source /opt/dtk/env.sh
COPY requirements.txt requirements.txt
COPY ../requirements.txt requirements.txt
RUN pip3 install -r requirements.txt
model.properties
View file @ 1707011e
# 模型唯一标识
modelCode=xxx
modelCode=479
# 模型名称
modelName=detr_pytorch
# 模型描述
modelDescription=DETR
modelDescription=DETR提出了一个十分简洁的目标检测pipeline, 即CNN提取基础特征, 送入Transformer做关系建模, 得到的输出通过二分图匹配算法与图片上的ground truth做匹配。
# 应用场景
appScenario=推理,训练,目标检测,网安,交通,政府
# 框架类型
......
requirements.txt
View file @ 1707011e
cython
pycocotools
submitit
pycocotools==2.0.7
submitit==1.5.1
torch>=1.5.0
torchvision>=0.6.0
scipy
onnx
onnxruntime
git+https://github.com/cocodataset/panopticapi.git
\ No newline at end of file
numpy
\ No newline at end of file
test.py
View file @ 1707011e
import os
import argparse
import os
import cv2
import torch
......@@ -8,16 +5,13 @@ import argparse
from datasets import build_dataset, get_coco_api_from_dataset
from models import build_test_model
from datasets.coco_eval import CocoEvaluator
import numpy as np
import util.misc as utils
def test_img(args, model, postprocessors, save_path):
dataset_test = build_dataset(
image_set="val", args=args, eval_in_training_set=False,
)
'''test img'''
dataset_test = build_dataset(image_set='val', args=args)
sampler_test = torch.utils.data.SequentialSampler(dataset_test)
data_loader_test = torch.utils.data.DataLoader(
dataset_test,
......@@ -34,7 +28,7 @@ def test_img(args, model, postprocessors, save_path):
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
outputs = model(img_data)
orig_target_sizes = torch.stack([t["orig_size"] for t in target], dim=0)
orig_target_sizes = torch.stack([t["orig_size"] for t in targets], dim=0)
result = postprocessors['bbox'](outputs, orig_target_sizes)
if 'segm' in postprocessors.keys():
target_sizes = torch.stack([t["size"] for t in targets], dim=0)
......@@ -73,7 +67,6 @@ def test_img(args, model, postprocessors, save_path):
coco_evaluator.summarize()
print(coco_evaluator)
def get_parser():
parser = argparse.ArgumentParser("DETR Detector", add_help=False)
parser.add_argument('--dataset_file', default='coco')
......@@ -81,8 +74,7 @@ def get_parser():
parser.add_argument('--coco_panoptic_path', type=str)
parser.add_argument('--remove_difficult', action='store_true')
parser.add_argument("--save_path", default="./result_img", type=str)
parser.add_argument('--output_dir', default='',
help='path where to save, empty for no saving')
parser.add_argument('--lr_backbone', default=1e-5, type=float)
parser.add_argument('--device', default='cuda',
help='device to use for training / testing')
parser.add_argument("--pre_trained_model", default="")
......@@ -116,17 +108,9 @@ def get_parser():
# Loss
parser.add_argument('--no_aux_loss', dest='aux_loss', action='store_false',
help="Disables auxiliary decoding losses (loss at each layer)")
# * eval technologies
parser.add_argument("--eval", action="store_true")
# eval in training set
parser.add_argument("--eval_in_training_set", default=False, action="store_true")
# topk for eval
parser.add_argument("--topk", default=100, type=int)
# * training technologies
parser.add_argument("--use_fp16", default=False, action="store_true")
parser.add_argument("--use_checkpoint", default=False, action="store_true")
return parser
if __name__ == "__main__":
args = get_parser().parse_args()
device = torch.device(args.device)
......@@ -138,7 +122,6 @@ if __name__ == "__main__":
model.to(device)
checkpoint = torch.load(args.pre_trained_model, map_location='cpu')
model.load_state_dict(checkpoint["model"], False)
model.eval()
test_img(args, model, postprocessors, args.save_path)
train.sh
View file @ 1707011e
#!/bin/bash
echo "Export params ..."
export HIP_VISIBLE_DEVICES=1,2,3,4 # 自行修改为训练的卡号和数量
export HIP_VISIBLE_DEVICES=0,1,2,3 # 自行修改为训练的卡号和数量
export HSA_FORCE_FINE_GRAIN_PCIE=1
export USE_MIOPEN_BATCHNORM=1
echo "Training start ..."
# coco_path是训练数据集地址，数据是coco format
# python -m torch.distributed.launch --nproc_per_node=4 --use_env main.py --coco_path /path/to/coco
python -m torch.distributed.launch --nproc_per_node=4 --use_env main.py --coco_path /home/datasets/COCO2017/images
python -m torch.distributed.launch --nproc_per_node=4 --use_env main.py --coco_path /path/to/coco
# python -m torch.distributed.launch --nproc_per_node=4 --use_env main.py --coco_path /home/datasets/COCO2017
echo "Training finished."
val.sh
View file @ 1707011e
#!/bin/bash
echo "Testing start ..."
export HIP_VISIBLE_DEVICES=1
export HSA_FORCE_FINE_GRAIN_PCIE=1
export USE_MIOPEN_BATCHNORM=1
# resume是待测试模型地址
# coco_path是训练数据集地址，数据是coco format
# python main.py --batch_size 2 --no_aux_loss --eval --resume path/of/model --coco_path /path/to/coco
# resume: 待测试模型地址
# coco_path: 训练数据集地址, 数据是coco format
python main.py --batch_size 2 --no_aux_loss --eval --resume /path/of/model --coco_path /path/of/coco_data
# python main.py --batch_size 2 --no_aux_loss --eval --resume https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth --coco_path /home/datasets/COCO2017
python main.py --batch_size 2 --no_aux_loss --eval --resume ./results/checkpoint.pth --coco_path /home/datasets/COCO2017
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment