Single channel image training？

[UNeedCryDear]

🚀 Feature

Thanks for this great work!
Can a parameter be set up to train and predict single channel image，like gray image.

Motivation

In the actual application process, many times do not need three channel BGR-image, but a single channel gray-image. I think a lot of people will have this situation. So I hope to have a parameter that can train and predict single channel images to speed up. From three channels to single channel, I think this speed should be nearly tripled.

Pitch

dataset=cv2.imread(train_data,0)
train(dataset,ch=1)
save(single_model)
detect_set=cv2.imread(detet_img,0)
model.predict(detect_set)

Alternatives

Additional context

[glenn-jocher]

@UNeedCryDear single-channel to 3-channel image you will see about zero speedup. 95% of FLOPS and parameter counts are not in the initial convolution. You can verify this by creating a 1-channel model in PyTorch Hub and comparing it's metrics against the default 3-channel model:

import torch

# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=False, autoshape=False, channels=3)
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=False, autoshape=False, channels=1)

See PyTorch Hub Tutorial to build your own single-channel models, though training single-channel would need some modifications to the dataloader.

YOLOv5 Tutorials

• Train Custom Data  🚀 RECOMMENDED
• Weights & Biases Logging  🌟 NEW
• Multi-GPU Training
• PyTorch Hub  ⭐ NEW
• ONNX and TorchScript Export
• Test-Time Augmentation (TTA)
• Model Ensembling
• Model Pruning/Sparsity
• Hyperparameter Evolution
• Transfer Learning with Frozen Layers  ⭐ NEW
• TensorRT Deployment

[UNeedCryDear]

@glenn-jocher Thank you for your reply.
In yolovs.onnx, there are three output with shape[bs,ch,h,w,num_classes].

output1: [1,3,80,80,85]

output2:[1,3,40,40,85]

output3:[1,3,20,20,85]
When deploying onnx_model in C + +, in the case of three channels, I had to traverse 25200（38080+34040+32020） times to judge each anchors_box(each length is 85). But if it is a single channel, I only need to traverse 8400（25200/3）times.

Because my images are grayscale. When converting to three channels, the data of each channel is the same, which will cause huge waste. So I'm going to ask if there's any way to get a single channel model.

Thank you again for your reply. I'll try it myself.

[glenn-jocher]

@UNeedCryDear your interpretation of the output shape being a function of the channel count is incorrect. Output shape is a function of anchor count and image size only, not channel count.

[UNeedCryDear]

@glenn-jocher Well, that should be my misunderstanding

[glenn-jocher]

@UNeedCryDear yes these are 3 anchors times 85 values (4 box coordinates, 1 objectness, 80 classes), at each of 3 output strides (8, 16, 32) for small, medium, large objects.

[Guemann-ui]

Hi ! What I understood from your discussion is that training a model with only one channel or 3 channels has no impact on model speed and inference! right?

[glenn-jocher]

@besmaGuesmi yes correct, input channel count has little effect on the model. You can create models with different input channels like this and compare stats and speeds:

import torch

for ch in 1, 3, 4, 10:
    model = torch.hub.load('ultralytics/yolov5', 'yolov5s', channels=ch)
    results = model(np.array((640, 640, ch)))
    results.print()

EDIT: Please git pull or git clone again to update your code as #5542 and #5543 fixed a bug in channel count adjustment with hub models. Above example works correctly with #5542 and #5543.

[spacewalk01]

how about single channel onnx model? I trained a model on gray scale images. I would like to export it to onnx with a single channel

[UNeedCryDear]

@spacewalk01

yolov5/export.py

Line 504 in 2f1eb21

im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection

change it to

im = torch.zeros(batch_size, 1, *imgsz).to(device)

and try again:
python detect.py --weights yours_best.pt --imgsz [640,640] --include onnx