YoloV8 WebGPU fixes (#8057)

* Bump up input size to 416, show if webgpu is not supported

* Minor fix in export_model
This commit is contained in:
Ahmed Harmouche
2024-12-05 16:23:45 +01:00
committed by GitHub
parent 78c01a5c2b
commit c6f5bb03fa
3 changed files with 13 additions and 6 deletions

View File

@@ -12,7 +12,7 @@ if __name__ == "__main__":
yolo_infer = YOLOv8(w=0.25, r=2.0, d=0.33, num_classes=80)
state_dict = safe_load(get_weights_location(yolo_variant))
load_state_dict(yolo_infer, state_dict)
prg, inp_sizes, out_sizes, state = export_model(yolo_infer, Device.DEFAULT.lower(), Tensor.randn(1,3,256,256))
prg, inp_sizes, out_sizes, state = export_model(yolo_infer, Device.DEFAULT.lower(), Tensor.randn(1,3,416,416))
dirname = Path(__file__).parent
safe_save(state, (dirname / "net.safetensors").as_posix())
with open(dirname / f"net.js", "w") as text_file:

View File

@@ -95,6 +95,7 @@
</head>
<body>
<h2>YOLOv8 tinygrad WebGPU</h2>
<h2 id="wgpu-error" style="display: none; color: red;">Error: WebGPU is not supported in this browser</h2>
<div class="video-container">
<video id="video" muted autoplay playsinline></video>
<canvas id="canvas"></canvas>
@@ -107,7 +108,7 @@
</div>
<script>
let net = null;
const modelInputSize = 256
const modelInputSize = 416;
let lastCalledTime;
let fps = 0, accumFps = 0, frameCounter = 0;
@@ -117,6 +118,7 @@
const offscreenCanvas = document.createElement('canvas');
const fpsMeter = document.getElementById('fps-meter');
const loadingContainer = document.getElementById('div-loading');
const wgpuError = document.getElementById('wgpu-error');
offscreenCanvas.width = modelInputSize;
offscreenCanvas.height = modelInputSize;
const offscreenContext = offscreenCanvas.getContext('2d');
@@ -147,7 +149,7 @@
lastCalledTime = now;
accumFps += 1/delta;
if (frameCounter++ >= 30) {
if (frameCounter++ >= 10) {
fps = accumFps/frameCounter;
frameCounter = 0;
accumFps = 0;
@@ -206,7 +208,12 @@
async function detectObjectsOnFrame(offscreenContext) {
if (!net) {
net = await loadNet(await getDevice());
let device = await getDevice();
if (!device) {
wgpuError.style.display = "block";
loadingContainer.style.display = "none";
}
net = await loadNet(device);
loadingContainer.style.display = "none";
}
let start = performance.now();
@@ -239,7 +246,7 @@
}
const getDevice = async () => {
if (!navigator.gpu) error("WebGPU not supported.");
if (!navigator.gpu) return false;
const adapter = await navigator.gpu.requestAdapter();
return await adapter.requestDevice();
};