python代码片段,红色部分的切片操作:
def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
labels=(), max_det=300):
nc = prediction.shape[2] - 5 # number of classes
xc = prediction[..., 4] > conf_thres # candidates
output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
print('-----------------------------------------')
for xi, x in enumerate(prediction): # image index, image inference
begin_time = datetime.datetime.now()
x = x[xc[xi]] # confidence
end_time = datetime.datetime.now()
d_time = end_time - begin_time
print(str(xi) + ': ' + str(d_time.microseconds))
其中prediction的形状是torch.Size([16, 15120, 85])
x = x[xc[xi]]这一句在循环中的执行时间如下,单位是微秒:
C++ Libtorch代码片段:
第一次forward的时间也比pytorch长的多,请教有没有人解决过这个问题?
我现在使用的libtorch相关的lib和dll都是从python那边拷贝过来的。