YoloV8 single channel train + Onnx trans

 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