yzqxy / Yolov8_obb_Prune_Track

GNU General Public License v3.0
181 stars 14 forks source link

training speed is very slow #20

Open XiaoJiNu opened 1 year ago

XiaoJiNu commented 1 year ago

Hi, I train the v8obb _n model with 300000 images and image size is 64, one epoch training time is 6 hours.

I use one A100 GPU and batch size is 128. The training time is 0.5 hour when I trained yolov5 obb with the same configuration.

what's the problem?how to fix.the bug?

cl886699 commented 1 year ago

replace

rotated_ious = []
for b1, b2 in zip(box1, box2):
    b1=b1.unsqueeze(0)
    b2=b2.unsqueeze(0)
    rotated_ious.append(box_iou_rotated(b1, b2).squeeze(0).squeeze(0))
return torch.stack(rotated_ious, axis=0)

with

return box_iou_rotated(box1, box2, aligned=True)
yzqxy commented 1 year ago

replace

rotated_ious = []
for b1, b2 in zip(box1, box2):
    b1=b1.unsqueeze(0)
    b2=b2.unsqueeze(0)
    rotated_ious.append(box_iou_rotated(b1, b2).squeeze(0).squeeze(0))
return torch.stack(rotated_ious, axis=0)

with

return box_iou_rotated(box1, box2, aligned=True)

This is a good way,thanks for your answer

XiaoJiNu commented 1 year ago

@cl886699 Thanks a lot

BuFanZHANG commented 12 months ago

replace

rotated_ious = []
for b1, b2 in zip(box1, box2):
    b1=b1.unsqueeze(0)
    b2=b2.unsqueeze(0)
    rotated_ious.append(box_iou_rotated(b1, b2).squeeze(0).squeeze(0))
return torch.stack(rotated_ious, axis=0)

with

return box_iou_rotated(box1, box2, aligned=True)

hello,i just want to know which code file I could replace it .thank you