补全yolov5的代码

2025-03-13 17:48:59 +08:00
parent 6daea23f0a
commit 9d369b9898
42 changed files with 8571 additions and 171 deletions
--- a/yolov5/train.py
+++ b/yolov5/train.py
@ -409,7 +409,7 @@ def train(hyp, opt, device, callbacks):
                    imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False)

            # Forward
-            with torch.amp.autocast(device_type='cuda', enabled=amp):
+            with torch.amp.autocast(device_type='cuda',enabled=amp):
                pred = model(imgs)  # forward
                loss, loss_items = compute_loss(pred, targets.to(device))  # loss scaled by batch_size
                if RANK != -1: