YOLOV5转换libtorch(GPU)标准代码

最近整理了一些转模型经常碰到的问题让我很是苦恼，在次整理给大家，希望大家少走弯路，早日项目娄底。以下都是在c++中调用libtorch模型时出错。

问题总汇

我们要怎么确定是不是我们的模型出错了呢？一般模型出错都是有窍门的。他一般都会伴随着C10这几个单词来的，要不就是torch::jit。

最常见也是最简单的报错

error while loading shared libraries: libtorch_cuda.so: cannot open shared object file: No such file or directory

这个保存原因就要看看你下载的libtorch是不是gpu版本的。在libtorch中有一个build-version。打开里面有libtorch+cu才是gpu版本的，才可以使用gpu。调用方法如下：

//m_modelPath为模型的路径
m_module = torch::jit::load(m_modelPath);
    //m_module.to(m_deviceType);
m_module.to(at::kCUDA);

找不到模型

terminate called after throwing an instance of ‘c10::Error’
what(): open file failed because of errno 2 on fopen: , file path:
这也是大家会经常碰到的错误，这里建议直接写绝对路径。

模型转换gpu版本不对

terminate called after throwing an instance of ‘c10::Error’
what(): isTuple() INTERNAL ASSERT FAILED at “/home/ls/sdk/libtorch/include/ATen/core/ivalue_inl.h”:931, please report a bug to PyTorch. Expected Tuple but got GenericList

这个错误很是恼火，我把底层翻了个地朝天，终于找到了。这里大概意思就是模型中没有使用Gpu.
这里我把我的转换代码贴出来，大家可以参考。

import argparse
import sys
import time

sys.path.append('./')  # to run '$ python *.py' files in subdirectories

import torch
import torch.nn as nn
#这里添加这个是防止找不到model路径
sys.path.append('/media/ls/5f831db7-8671-43b0-8d53-0ad743c608e1/3Ddedict/axleDet_yolov5_dianyun5')
sys.path.append('/media/ls/5f831db7-8671-43b0-8d53-0ad743c608e1/3Ddedict/axleDet_yolov5_dianyun5/models')
import models
from models.experimental import attempt_load
from utils.activations import Hardswish, SiLU
from utils.general import set_logging, check_img_size
from utils.torch_utils import select_device
#输出torch是否支持cuda。输出1为支持，0为不支持。
print(torch.cuda.is_available())
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='需要转换的模型路径', help='weights path')  # from yolov5/models/
    #图片大小很重要 要和部署时的大小一致
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size')  # height, width
    #这里默认为1
    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
    #生成onnx模型，默认生成。
    parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes')
    parser.add_argument('--grid', action='store_true', help='export Detect() layer grid')
    #官方给给出的控制gpu，cpu的代码段。问题就出现在这里，所以为直接注释，具体怎么去调用gpu，大家往下看，我写了个最简单也是最不容易出错的。
    #parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    opt = parser.parse_args()
    opt.img_size *= 2 if len(opt.img_size) == 1 else 1  # expand
    print(opt)
    set_logging()
    t = time.time()
    #这里，如果检测到有cuda就使用cuda否则使用cpu
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(device)
    # Load PyTorch model
    #device = select_device(opt.device)
    #map_location=torch.device('cuda') 这里直接使用cuda进行前向传播，这是一定不会出错的。
    model = attempt_load(opt.weights, map_location=torch.device('cuda'))  # load FP32 model
    labels = model.names

    # Checks
    gs = int(max(model.stride))  # grid size (max stride)
    opt.img_size = [check_img_size(x, gs) for x in opt.img_size]  # verify img_size are gs-multiples

    # Input
    #把图片直接放入cuda进行推理，减少犯错可能。
    img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device='cuda')  # image size(1,3,320,192) iDetection

    # Update model
    for k, m in model.named_modules():
        m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility
        if isinstance(m, models.common.Conv):  # assign export-friendly activations
            if isinstance(m.act, nn.Hardswish):
                m.act = Hardswish()
            elif isinstance(m.act, nn.SiLU):
                m.act = SiLU()
        # elif isinstance(m, models.yolo.Detect):
        #     m.forward = m.forward_export  # assign forward (optional)
    #model.model[-1].export = not opt.grid  # set Detect() layer grid export
    model.model[-1].export = False
    y = model(img)  # dry run

    # TorchScript export
    try:
    	
        print('\nStarting TorchScript export with torch %s...' % torch.__version__)
       #你要起的给生成模型起的名字。
        f = opt.weights.replace('.pt', '.torchscript.pt')  # filename
        ts = torch.jit.trace(model, img)
        ts.save(f)
        print('TorchScript export success, saved as %s' % f)
    except Exception as e:
        print('TorchScript export failure: %s' % e)

    # ONNX export
    try:
        import onnx

        print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
        f = opt.weights.replace('.pt', '.onnx')  # filename
        torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
                          output_names=['classes', 'boxes'] if y is None else ['output'],
                          dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # size(1,3,640,640)
                                        'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)

        # Checks
        onnx_model = onnx.load(f)  # load onnx model
        onnx.checker.check_model(onnx_model)  # check onnx model
        # print(onnx.helper.printable_graph(onnx_model.graph))  # print a human readable model
        print('ONNX export success, saved as %s' % f)
    except Exception as e:
        print('ONNX export failure: %s' % e)

    # CoreML export
    try:
        import coremltools as ct

        print('\nStarting CoreML export with coremltools %s...' % ct.__version__)
        # convert model from torchscript and apply pixel scaling as per detect.py
        model = ct.convert(ts, inputs=[ct.ImageType(name='image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
        f = opt.weights.replace('.pt', '.mlmodel')  # filename
        model.save(f)
        print('CoreML export success, saved as %s' % f)
    except Exception as e:
        print('CoreML export failure: %s' % e)

    # Finish
    print('\nExport complete (%.2fs). Visualize with https://github.com/lutzroeder/netron.' % (time.time() - t))