增加跟踪的一些逻辑代码

2025-03-24 17:42:10 +08:00
parent bba39adcfc
commit f32cd5b9a2
47 changed files with 3395 additions and 17 deletions
--- a/yolov5/models/common.py
+++ b/yolov5/models/common.py
@ -891,7 +891,7 @@ class AutoShape(nn.Module):
            x = np.ascontiguousarray(np.array(x).transpose((0, 3, 1, 2)))  # stack and BHWC to BCHW
            x = torch.from_numpy(x).to(p.device).type_as(p) / 255  # uint8 to fp16/32

-        with amp.autocast(autocast):
+        with torch.amp.autocast(device_type='cuda', enabled=autocast):
            # Inference
            with dt[1]:
                y = self.model(x, augment=augment)  # forward