跳转至

训练AlphaPose

训练代码

https://github.com/wojiazaiyugang/AlphaPose

训练环境

docker pull wojiazaiyugang/alpha-pose:latest

进入AlphaPose文件夹，执行python setup.py build develop

数据集准备

训练时使用COCO格式的数据集。首先参考使用labelme进行关键点的标注 标注完成后使用generate_coco_from_labelme_human_key_point.py脚本生成COCO格式的数据集。

配置文件

DATASET:
  TRAIN:
    TYPE: 'Mscoco'
    ROOT: './data/coco/'
    IMG_PREFIX: 'train2017'
    ANN: 'annotations/person_keypoints_train2017.json'
    AUG:
      FLIP: true
      ROT_FACTOR: 40
      SCALE_FACTOR: 0.3
      NUM_JOINTS_HALF_BODY: 8
      PROB_HALF_BODY: -1
  VAL:
    TYPE: 'Mscoco'
    ROOT: './data/coco/'
    IMG_PREFIX: 'val2017'
    ANN: 'annotations/person_keypoints_val2017.json'
  TEST:
    TYPE: 'Mscoco_det'
    ROOT: './data/coco/'
    IMG_PREFIX: 'val2017'
    DET_FILE: './exp/json/test_det_yolo.json'
    ANN: 'annotations/person_keypoints_val2017.json'
DATA_PRESET:
  TYPE: 'simple'
  SIGMA: 2
  NUM_JOINTS: 17
  IMAGE_SIZE:
  - 256
  - 192
  HEATMAP_SIZE:
  - 64
  - 48
MODEL:
  TYPE: 'FastPose'
  PRETRAINED: '' # 载入预训练模型
  TRY_LOAD: ''
  NUM_DECONV_FILTERS:
  - 256
  - 256
  - 256
  NUM_LAYERS: 50
LOSS:
  TYPE: 'MSELoss'
DETECTOR:
  NAME: 'yolo'
  CONFIG: 'detector/yolo/cfg/yolov3-spp.cfg'
  WEIGHTS: 'detector/yolo/data/yolov3-spp.weights' # 验证的时候会有一个rcnn map，是用yolov3检测bbox然后检测关键点的，这个权重文件需要提前下载放到detector/yolo/data里
  NMS_THRES: 0.6
  CONFIDENCE: 0.05
TRAIN:
  WORLD_SIZE: 4
  BATCH_SIZE: 32
  BEGIN_EPOCH: 0
  END_EPOCH: 200
  OPTIMIZER: 'adam'
  LR: 0.001
  LR_FACTOR: 0.1
  LR_STEP:
  - 90
  - 120
  DPG_MILESTONE: 140
  DPG_STEP:
  - 160
  - 190

从头训练

python scripts/train.py --exp-id boxer_key_point --cfg configs/coco/resnet/256x192_res50_lr1e-3_1x.yaml

batch=192时大约占用30G显存

载入预训练模型

配置文件中MODEL-PRETRAINED配置要加载的模型即可

模型转onnx

python scripts/pytorch2onnx.py --cfg /workspace/AlphaPose/configs/coco/resnet/256x192_res50_lr1e-3_1x.yaml --checkpoint exp/boxer_key_point-256x192_res50_lr1e-3_1x.yaml/model_127.pth --batch_size 1

评论