torchscript 和 C++ 接口中的 Wrong input shape 错误

Question

我正在尝试连接 libtorch 和 OpenCV，以使用 Yolov5 模型预测类。我使用的权重是yolov5s.pt。源代码是

cv::Mat image = file->input_image(); // read image and resize into 640x640

auto tensor = torch::from_blob(image.data, {image.rows,image.cols,3}, torch::kFloat);
tensor = tensor.view({1,640,640,3});
std::cout << tensor.sizes() << std::endl;

try {
  auto output = model.forward({tensor}).toTensor();
  std::cout << output.sizes() << std::endl;
} catch (std::runtime_error & e) {
  std::cerr << "[X] Error: " << e.what() << std::endl;
  return;
}

错误信息

RuntimeError: Given groups=1, weight of size [32, 3, 6, 6], expected input[1, 640, 640, 3] to have 3 
channels, but got 640 channels instead

回溯

Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/models/yolo.py", line 59, in forward
    model23 = self.model
    _0 = getattr(model23, "0")
    _25 = (_2).forward((_1).forward((_0).forward(x, ), ), )
                                     ~~~~~~~~~~~ <--- HERE
    _26 = (_4).forward((_3).forward(_25, ), )
    _27 = (_6).forward((_5).forward(_26, ), )
  File "code/__torch__/models/common.py", line 12, in forward
    act = self.act
    conv = self.conv
    _0 = (act).forward((conv).forward(x, ), )
                        ~~~~~~~~~~~~~ <--- HERE
    return _0
class C3(Module):
  File "code/__torch__/torch/nn/modules/conv.py", line 12, in forward
    bias = self.bias
    weight = self.weight
    x0 = torch._convolution(x, weight, bias, [2, 2], [2, 2], [1, 1], False, [0, 0], 1, False, False, True, True)
         ~~~~~~~~~~~~~~~~~~ <--- HERE
    return x0

Answer 1

解决方案非常简单。我觉得太尴尬了，但你不应该。

这是解决方案

// forgot to add these both lines
// the yolov5 is expects [BATCH, CHANNEL, WIDTH, HEIGHT]
tensor = tensor.permute({2,0,1});
tensor = tensor = tensor.unsqueeze(0);

std::cout << tensor.sizes() << std::endl;

torchscript 和 C++ 接口中的 Wrong input shape 错误

Wrong input shape error in torchscript and C++ interface

c++

opencv

libtorch

yolov5