training speed is very slow

[XiaoJiNu]

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]

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]

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]

@cl886699 Thanks a lot

[BuFanZHANG]

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