yolov8目前不支持单通道图片训练,需要修改后才能支持。本文将介绍如何修改yolov8代码,来训练单通道图的yolov8模型,以及使用onnx进行模型转换的简单实践。
1、修改代码
git diff ultralytics/utils/checks.py
diff --git a/ultralytics/utils/checks.py b/ultralytics/utils/checks.py
index 1ab031ba..c8fadcdc 100644
--- a/ultralytics/utils/checks.py
+++ b/ultralytics/utils/checks.py
@@ -648,7 +648,7 @@ def check_amp(model):
try:
from ultralytics import YOLO
- assert amp_allclose(YOLO("yolov8n.pt"), im)
+ # assert amp_allclose(YOLO("yolov8n.pt"), im)
LOGGER.info(f"{prefix}checks passed ✅")
except ConnectionError:
LOGGER.warning(f"{prefix}checks skipped ⚠️, offline and unable to download YOLOv8n. {warning_msg}")
diff --git a/ultralytics/nn/tasks.py b/ultralytics/nn/tasks.py
index 64ee7f50..c8f37677 100644
--- a/ultralytics/nn/tasks.py
+++ b/ultralytics/nn/tasks.py
@@ -263,6 +264,8 @@ class BaseModel(nn.Module):
if not hasattr(self, "criterion"):
self.criterion = self.init_criterion()
+ # import ipdb;ipdb.set_trace()
+ batch['img'] = batch['img'][:,:1,:,:]
preds = self.forward(batch["img"]) if preds is None else preds
return self.criterion(preds, batch)
diff --git a/ultralytics/nn/autobackend.py b/ultralytics/nn/autobackend.py
index 4d8c69c5..98469690 100644
--- a/ultralytics/nn/autobackend.py
+++ b/ultralytics/nn/autobackend.py
@@ -502,7 +502,7 @@ class AutoBackend(nn.Module):
"""
return torch.tensor(x).to(self.device) if isinstance(x, np.ndarray) else x
- def warmup(self, imgsz=(1, 3, 640, 640)):
+ def warmup(self, imgsz=(1, 1, 640, 640)):
"""
Warm up the model by running one forward pass with a dummy input.
diff --git a/ultralytics/engine/validator.py b/ultralytics/engine/validator.py
index 41be54c1..82ed2f03 100644
--- a/ultralytics/engine/validator.py
+++ b/ultralytics/engine/validator.py
@@ -154,7 +154,8 @@ class BaseValidator:
self.dataloader = self.dataloader or self.get_dataloader(self.data.get(self.args.split), self.args.batch)
model.eval()
- model.warmup(imgsz=(1 if pt else self.args.batch, 3, imgsz, imgsz)) # warmup
+ model.warmup(imgsz=(1 if pt else self.args.batch, 1, imgsz, imgsz)) # warmup gray image
+ # model.warmup(imgsz=(1 if pt else self.args.batch, 3, imgsz, imgsz)) # warmup
self.run_callbacks("on_val_start")
dt = (
@@ -175,6 +176,7 @@ class BaseValidator:
# Inference
with dt[1]:
+ batch["img"] = batch["img"][:,:1,:,:]
preds = model(batch["img"], augment=augment)
# Loss
diff --git a/ultralytics/engine/predictor.py b/ultralytics/engine/predictor.py
index e925902f..32745d6f 100644
--- a/ultralytics/engine/predictor.py
+++ b/ultralytics/engine/predictor.py
@@ -261,7 +261,8 @@ class BasePredictor:
# Warmup model
if not self.done_warmup:
- self.model.warmup(imgsz=(1 if self.model.pt or self.model.triton else self.dataset.bs, 3, *self.imgsz))
+ self.model.warmup(imgsz=(1 if self.model.pt or self.model.triton else self.dataset.bs, 1, *self.imgsz))
+ # self.model.warmup(imgsz=(1 if self.model.pt or self.model.triton else the dataset.bs, 3, *self.imgsz))
self.done_warmup = True
self.seen, self.windows, self.batch = 0, [], None
2、修改配置文件
ultralytics/cfg/models/v8/yolov8-xxx.yaml
添加ch: 1
3、onnx模型转换
ultralytics/engine/exporter.py
修改im = torch.zeros(self.args.batch, 3, *self.imgsz).to(self.device)
改为im = torch.zeros(self.args.batch, 1, *self.imgsz).to(self.device)
4、测试pt模型
找到对应任务的ultralytics/models/yolo/classify(自己任务)/predict.py
添加:img = img[:,:1,:,:]
参考
SingleChannel Train · Issue #7526 · ultralytics/ultralytics · GitHub