yolov5_obb的兼容性改进
yolov5_obb的兼容性改进
最近在学习yolo的旋转框检测,选用了yolov5_obb进行测试,结果代码由于各种问题跑不起来。我将里面用到的废弃函数进行了更改,以便能用用新环境进行运行。
我将修改完的代码上传到了我的Github仓库里。
改进方法
utils/nms_rotated/nms_rotated_wrapper.py
运行的时候会显示too_small
和inds
这两个变量和ori_inds
不在同一个设备上,需要添加代码将两个变量移动到ori_inds
所在的设备上。
1 | ori_inds = torch.arange(dets_th.size(0)) # 0 ~ n-1 |
utils/nms_rotated/src/poly_nms_cuda.cu
THC库已经被弃用,我们需要用ATen库进行更换
这里我是根据这个博客进行更改的
删除#include <THC/THC.h>
头文件,把所有的THCudaCheck(cudaGetLastError());
替换成 AT_CUDA_CHECK(cudaGetLastError());
对每个.cu文件中调用THCCeilDiv(x,y)
的地方,都把这个函数换成 (x+y-1)/y
的形式,如:
1 | // const int col_blocks = THCCeilDiv(n_polys, threadsPerBlock); |
添加头文件#include <ATen/cuda/ThrustAllocator.h>
注释掉THCState *state = at::globalContext().lazyInitCUDA();
修改THCudaMalloc
函数:
1 | // mask_dev = (unsigned long long*) THCudaMalloc(state, boxes_num * col_blocks * sizeof(unsigned long long)); |
替换THCudaFree
1 | // THCudaFree(state, mask_dev); |
将THCudaCheck
替换成C10_CUDA_CHECK
将AT_CHECK
替换为TORCH_CHECK
utils/datasets.py
将所有的np.int
改写为int
utils/general.py
torch.load
函数添加weights_only=False
参数
utils/loss.py
这个文件会报错说数据类型转换错误。需要先将变量转换为long,注意这里不是原地转换,需要用赋值的形式进行转换。
1 | # indices.append((b, a, gj.clamp_(0, feature_wh[1] - 1), gi.clamp_(0, feature_wh[0] - 1))) # image, anchor, grid indices |
utils/plots.py
self.font.getsize
函数已经被废弃。
1 | # w, h = self.font.getsize(label) # text width, height |
models/experimental.py
torch.load
函数添加weights_only=False
参数
train.py
torch.load
函数添加weights_only=False
参数, 注释掉下面两行防止更新numpy。
1 | # check_git_status() |
编译C++代码
原作者用C++写了一些函数,需要编译后才能使用。但是我按照他写的命令缺一直报错说ModuleNotFoundError: No module named 'torch'
,显然我的环境里是有torch的,但是编译的时候报错说没有。
非常玄学的是,我确实有一个环境能够通过python setup.py develop
命令来进行编译的,我记得当时我安装了一个pycocotools库之后就可以了。但是过了一天,我想从头到尾再创建一个新的虚拟环境过一遍的时候发现又不能编译了,我查了一下pycocotools仅仅是python对coco数据集做的小工具,应该不是我编译成功的关键要素。我也尝试过sudo apt install python-dev
安装编译工具,但是没有用。我导出pip requirement和conda list并且仔细对比了一下发现里面的东西都是一样的。我寻找了很多网页都没有发现是怎么回事,有一些人和我一样遇到了一样的问题,但是都没有解决方案。
我觉得可能是代码与新环境不匹配的问题。我在和chatgpt讨论如何更改代码的时候,他提到了一个新的执行命令:
1 | python setup.py build_ext --inplace |
我发现这个命令能够成功编译,而且运行训练任务测试能够完美运行。我不知道这算不算解决了问题,我不知道解决问题的原因和答案,换了一个命令就成功了。
PS:封面图来源:天音かなた💫ホロライブ