@@ -41,11 +41,23 @@ def test_yolov8_loss_wrapper():
4141 x = torch .randn ((batch_size , 3 , 640 , 640 )) # YOLOv8 expects (B, 3, 640, 640)
4242
4343 # Create targets
44- targets = []
44+ """ targets = []
4545 for _ in range(batch_size):
4646 boxes = torch.tensor([[0.1, 0.1, 0.3, 0.3], [0.5, 0.5, 0.8, 0.8]]) # [x1, y1, x2, y2]
4747 labels = torch.zeros(2, dtype=torch.long) # Use class 0 for testing
48- targets .append ({"boxes" : boxes , "labels" : labels })
48+ targets.append({"boxes": boxes, "labels": labels})"""
49+ targets = torch .tensor (
50+ [
51+ [0.0000 , 20.0000 , 0.7738 , 0.3919 , 0.4525 , 0.7582 ],
52+ [0.0000 , 20.0000 , 0.2487 , 0.4062 , 0.4966 , 0.5787 ],
53+ [0.0000 , 20.0000 , 0.5667 , 0.2772 , 0.0791 , 0.2313 ],
54+ [0.0000 , 20.0000 , 0.1009 , 0.1955 , 0.2002 , 0.0835 ],
55+ [1.0000 , 20.0000 , 0.7738 , 0.3919 , 0.4525 , 0.7582 ],
56+ [1.0000 , 20.0000 , 0.2487 , 0.4062 , 0.4966 , 0.5787 ],
57+ [1.0000 , 20.0000 , 0.5667 , 0.2772 , 0.0791 , 0.2313 ],
58+ [1.0000 , 20.0000 , 0.1009 , 0.1955 , 0.2002 , 0.0835 ],
59+ ]
60+ )
4961
5062 # Test training mode
5163 losses = wrapper (x , targets )
@@ -94,11 +106,23 @@ def test_yolov10_loss_wrapper():
94106 x = torch .randn ((batch_size , 3 , 640 , 640 )) # Standard YOLO input size
95107
96108 # Create targets
97- targets = []
109+ """ targets = []
98110 for _ in range(batch_size):
99111 boxes = torch.tensor([[0.1, 0.1, 0.3, 0.3], [0.5, 0.5, 0.8, 0.8]]) # [x1, y1, x2, y2]
100112 labels = torch.zeros(2, dtype=torch.long) # Use class 0 for testing
101- targets .append ({"boxes" : boxes , "labels" : labels })
113+ targets.append({"boxes": boxes, "labels": labels})"""
114+ targets = torch .tensor (
115+ [
116+ [0.0000 , 20.0000 , 0.7738 , 0.3919 , 0.4525 , 0.7582 ],
117+ [0.0000 , 20.0000 , 0.2487 , 0.4062 , 0.4966 , 0.5787 ],
118+ [0.0000 , 20.0000 , 0.5667 , 0.2772 , 0.0791 , 0.2313 ],
119+ [0.0000 , 20.0000 , 0.1009 , 0.1955 , 0.2002 , 0.0835 ],
120+ [1.0000 , 20.0000 , 0.7738 , 0.3919 , 0.4525 , 0.7582 ],
121+ [1.0000 , 20.0000 , 0.2487 , 0.4062 , 0.4966 , 0.5787 ],
122+ [1.0000 , 20.0000 , 0.5667 , 0.2772 , 0.0791 , 0.2313 ],
123+ [1.0000 , 20.0000 , 0.1009 , 0.1955 , 0.2002 , 0.0835 ],
124+ ]
125+ )
102126
103127 # Test training mode
104128 losses = wrapper (x , targets )
@@ -219,7 +243,7 @@ def loss(self, items):
219243 wrapper .train ()
220244 # Dummy input and targets
221245 x = torch .zeros ((1 , 3 , 416 , 416 ))
222- targets = [{ "boxes" : torch .zeros (( 1 , 4 )), "labels" : torch . zeros (( 1 ,))}]
246+ targets = torch .tensor ([[ 0.0000 , 20.0000 , 0.7738 , 0.3919 , 0.4525 , 0.7582 ]])
223247 losses = wrapper (x , targets )
224248 assert set (losses .keys ()) == {"loss_total" , "loss_box" , "loss_cls" , "loss_dfl" }
225249 assert losses ["loss_total" ].item () == 6.0 # sum([1.0, 2.0, 3.0])
@@ -264,7 +288,7 @@ def loss(self, items):
264288 wrapper = PyTorchYoloLossWrapper (test_model , name = "yolov8n" )
265289 wrapper .train ()
266290 x = torch .zeros ((1 , 3 , 416 , 416 ))
267- targets = [{ "boxes" : torch .zeros (( 1 , 4 )), "labels" : torch . zeros (( 1 ,))}]
291+ targets = torch .tensor ([[ 0.0000 , 20.0000 , 0.7738 , 0.3919 , 0.4525 , 0.7582 ]])
268292 losses = wrapper (x , targets )
269293 assert set (losses .keys ()) == {"loss_total" , "loss_box" , "loss_cls" , "loss_dfl" }
270294 assert losses ["loss_total" ].item () == 6.0
@@ -439,9 +463,7 @@ def loss(self, items):
439463 for batch_size in batch_sizes :
440464 for box_count in box_counts :
441465 x = torch .zeros ((batch_size , 3 , 416 , 416 ))
442- targets = [
443- {"boxes" : torch .zeros ((box_count , 4 )), "labels" : torch .zeros (box_count )} for _ in range (batch_size )
444- ]
466+ targets = torch .tensor ([[0.0000 , 20.0000 , 0.7738 , 0.3919 , 0.4525 , 0.7582 ]] * batch_size )
445467 losses = wrapper (x , targets )
446468
447469 # Verify loss structure
0 commit comments