CenterNet踩坑记

借鉴了两位大佬的训练指导(非常感谢)
https://blog.csdn.net/weixin_42634342/article/details/97756458
https://blog.csdn.net/weixin_41765699/article/details/100118353
记录一下自己实验的情况(纯小白)
我的数据集一共6069张,3885张训练,1213张验证,971张测试(其实用7:2:1划分即可,我分的太乱了)

首先是常规流程

我的环境为cuda10.0 cudnn7.6 pytorch1.1.0
装CenterNet环境结合官方和这个issue安装,切记2080显卡pytorch需要装1.1版本:
https://github.com/xingyizhou/CenterNet
https://github.com/xingyizhou/CenterNet/issues/356
重点:

训练自己的数据

1.在CenterNet/src/lib/datasets/dataset文件夹中,依据coco.py新建自己数据集的文件work.py

import torch.utils.data as data

class Work(data.Dataset):
  num_classes = 3
  default_resolution = [512, 512]
  mean = np.array([0.423766, 0.424928, 0.390202],
                   dtype=np.float32).reshape(1,