- 原repos地址为:https://github.com/princeton-vl/CornerNet
- 为了在服务器的cuda10电脑上用自己的数据集跑这份代码,主要有三点:
- 修改数据集接口
- 适应cuda10和pytorch1.5
- 训练配置
- 接下来就分点详述:
修改数据集接口:
- 有四个文件要修改:
- models/CornerNet.py:
- 72行model类的outdim变量改为自己数据集的类别数,如果类别加上背景为n+1,那么改为n
- config/CornerNet.json
- “train_split” 和 “val_split”(这里改为什么后面会提到)
- "categories"改为自己数据集的类别数,也是,如果加上背景是n+1那么改为n
- db/detection.py
- self._configs[“categories”] 设为自己的数据集的类别数,同上
- db/coco.py
- 这个是根据自己的数据集的路径灵活修改的,我以我的为例子,我先将我的数据集的标注转为coco格式的数据集,也就是有两个json文件分别是训练集train.json和测试集的标注test.json(根据我的需要我没有设验证集,有需要的灵活应变),此外还有所有的图片文件无论是训练集还是测试集全放在同一个文件夹中,总的来说,repos目录下的data文件夹中,有两个文件夹,一个是cocoapi所在文件夹,一个是vhr(我的数据集),vhr文件夹中有两个文件夹,images和annotations,images里面是所有的图片,annotations里面是两个json文件
- 根据上述路径,需要修改的是coco.py里MSCOCO类里面的几个变量:
- self._dataset:这个变量其实是一个字符串,原来的写法是为了,利用config 的json文件里面的train_split和val_split作为键从字典中取出对应路径的字符串,这个字符串对应着你标注文件的前缀,由于我的标注文件是train.json和test.json,所以我把self._dataset设为
同时将前面提到的CornerNet.json文件里面的两项设为:self._dataset = { "train": "train", "test": "test" }[self._split]
"train_split": "train", "val_split": "test",
- self._coco_dir :这个变量是为了指出数据集文件夹的名称,为了配合后续的label_dir和image_dir而使用的,因此我设为:
self._coco_dir = os.path.join(data_dir, "vhr")
- self._label_dir 和 self._label_file:label_dir是为了配合label_file指出json文件的路径和文件名的,只要最终label_flie这个变量指向你json文件即可,我设为:
self._label_dir = os.path.join(self._coco_dir, "annotations") self._label_file = os.path.join(self._label_dir, "{}.json") self._label_file = self._label_file.format(self._dataset)
- self._image_dir 这个变量最终要指向图片所在的文件夹的路径,所以我设为:
self._image_dir = os.path.join(self._coco_dir, "images")
- self._cat_ids 如果你的标注的cat_id是0到n-1的不间断的数字,那么这里设为
[i+1 for i in range(n)]
即可,其中n要替换为你的类别数,自此数据集接口修改大功告成
- self._dataset:这个变量其实是一个字符串,原来的写法是为了,利用config 的json文件里面的train_split和val_split作为键从字典中取出对应路径的字符串,这个字符串对应着你标注文件的前缀,由于我的标注文件是train.json和test.json,所以我把self._dataset设为
- models/CornerNet.py:
适应cuda10和pytorch1.5:
- 这个我找了很久才找到的方法,即使你按照repos给的环境安装指引去安装,只要你GPU的驱动依赖的是cuda10,就会出错,而且还不会告诉你错在哪里,就是训练卡了半小时最后告诉你unhandled cuda error,所以我在原来的环境的基础上,用这句命令行把pytorch从0.4更新为最新的1.5同时用cuda10:
conda install pytorch torchvision cudatoolkit=10.1 -c pytorch
并且重新编译了相关环境 - 但是这样是会出错的,原因是它的cptool本来就是写在pytorch0.4和cuda8上面的,所以要改cptool,简单的方法就是把cornernet-lite这个repos里写在pytorch1.0上的cptool直接覆盖掉原来的cptool,在这里添加链接描述,把新代码的_cptools文件夹覆盖到原来的_cptools文件夹上就行了
- 但是这样会出现新的问题,它不停地报warning,而且按网上说的方法去组织python报warning并没有用。warning的原因是新的_cptools是在pytorch1.0上写的,但是pytorch1.0以后已经抛弃了1.0中用kByte作为mask的做法而需要改为kBool,具体是两步修改:
- 把_cptools中*_pool.cpp,也就是说4个都要,中的kByte改为kBool
- 把sample/coco.py中77行的tagmasks改为np.tools
tag_masks = np.zeros((batch_size, max_tag_len), dtype=np.bool)
- 这样就大功告成了
其他配置
- 最后再讲一下其它配置,本来不需要说但是有一个地方有点坑,CornerNet.json里面的chunk_size看起来不明所以不知道是什么其实指的是每张卡上的batchsize,可以看出原来是在10张卡上训练的,但是多数人没有10张卡,所以会出错,我有4张卡,所以是:
"chunk_sizes": [4, 4, 4, 4],
,同时要把batch_size改为chunk_sizes的加和,我是16 - 此外还有就是如果是在服务器上挂后台训练并且把log输出到文件中,由于它用了tqdm,不够友好,所以我把train.py里面的tqdm去掉了