Win10下 pytorch模型在C++上的部署

Win10下 pytorch模型在C++上的部署 <一、>
一、使用python train好需要的模型（.pth）
二、将.pth转换为.pt文件
三、配置VS2017+libtorch+opencv,导入.pt模型进行编码处理

本次实验所需环境、IDE依赖如下：
python:pycharm+pytorch1.2.0;
C++:VS2017+libtorch1.2.0+opencv3.4.6;
win10+1050Ti

一、使用python train好需要的模型（.pth）
本次实验通过libtorch在C++上部署pytorch_yolov4，yolov4源代码来源：https://github.com/bubbliiiing/yolov4-pytorch，相关配置环境链接里有介绍，使用教学可以看B站bubbliiiing讲解视频，不在多做赘述。
二、将.pth转换为.pt文件
将pth转为pt的方法，libtorch有相关介绍，参照一下链接：https://pytorch.org/tutorials/advanced/cpp_export.html?highlight=libtorch，
我的方法如下：

在这里插入代码片
# -------------------------------------#
#       对数据集进行训练
# -------------------------------------#
import os
import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.optim as optim
from torch.utils.data import DataLoader
from tqdm import tqdm
import cv2
from PIL import Image
from torchvision import transforms
from nets.yolo4 import YoloBody
from nets.yolo_training import LossHistory, YOLOLoss, weights_init
from utils.dataloader import YoloDataset, yolo_dataset_collate

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 这里是自己的GPU编号，我是"0",原文是"8".


# ---------------------------------------------------#
#   获得类和先验框
# ---------------------------------------------------#
def get_classes(classes_path):
    '''loads the classes'''
    with open(classes_path) as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]
    return class_names


def get_anchors(anchors_path):
    '''loads the anchors from a file'''
    with open(anchors_path) as f:
        anchors = f.readline()
    anchors = [float(x) for x in anchors.split(',')]
    return np.array(anchors).reshape([-1, 3, 2])[::-1, :, :]


def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']
if __name__ == "__main__":
    #   是否使用Tensorboard
    Tensorboard = False
    #   是否使用Cuda
    #   没有GPU可以设置成False
    Cuda = False
    #   是否对损失进行归一化，用于改变loss的大小
    #   用于决定计算最终loss是除上batch_size还是除上正样本数量
    normalize = False
    input_shape = (416, 416)
    anchors_path = 'model_data/yolo_anchors.txt'
    classes_path = 'model_data/hotspot.txt'#类文件
    mosaic = False
    Cosine_lr = False  # 余弦退火衰减
    smoooth_label = 0
    class_names = get_classes(classes_path)
    anchors = get_anchors(anchors_path)
    num_classes = len(class_names)
    model = YoloBody(len(anchors[0]), num_classes)
    device = torch.device('cpu')
    model_path = ".../xxx.pth"

    print('Loading weights into state dict...')
    model_dict = model.state_dict()
    pretrained_dict = torch.load(model_path, map_location=torch.device("cpu"))
    pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
    model_dict.update(pretrained_dict)
    model.load_state_dict(model_dict)
    print('Finished!')

    # 生成pt模型的核心模块
    model = model.to(torch.device("cpu"))
    model.eval()
    var = torch.ones((1, 3, 416, 416))
    traced_script_module = torch.jit.trace(model, var)
    #traced_script_module.save("输出路径/.pt")

其实也就是将train.py的执行代码稍微做修改，只需要做一次输入，让模型执行一次，同时 torch.jit.trace（）会将这个过程trace下来，然后保存即可。
end!