1. ultralytics/cfg/models/v8/yolov8-xxx.yaml
添加ch: 1
2. ultralytics/utils/checks.py
注释assert amp_allclose(YOLO("yolov8n.pt"), im)
3. ultralytics/nn/tasks.py
添加:batch['img'] = batch['img'][:,:1,:,:]
4. ultralytics/nn/autobackend.py
修改:def warmup(self, imgsz=(1, 1, 640, 640)):
5. ultralytics/engine/validator.py
修改:model.warmup(imgsz=(1 if pt else self.args.batch, 1, imgsz, imgsz))
添加:batch["img"] = batch["img"][:,:1,:,:]
6. ultralytics/engine/predictor.py
修改:self.model.warmup(imgsz=(1 if self.model.pt or self.model.triton else self.dataset.bs, 1, *self.imgsz))
训练结束后如何转出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)
训练结束后如何测试pt模型
找到对应任务的ultralytics/models/yolo/classify(自己任务)/predict.py
添加:img = img[:,:1,:,:]
参考
github issue