Implementación del modelo de Pytorch ------- registro completo de implementación de Pytorch

Registro completo de la implementación del modelo Pytorch

guardar modelo de pytorch

Guardar modelo

torch.save(model, ‘net.pth’) # 保存整个神经网络的模型结构以及参数
torch.save(model, ‘net.pkl’) # 保存整个神经网络的模型结构以及参数
torch.save(model.state_dict(), ‘net_params.pth’) # 只保存模型参数
torch.save(model.state_dict(), ‘net_params.pkl’) # 只保存模型参数

Carga del modelo

model = torch.load(‘net.pth’) # 加载整个神经网络的模型结构以及参数
model = torch.load(‘net.pkl’) # 加载整个神经网络的模型结构以及参数
model.load_state_dict(torch.load(‘net_params.pth’)) # 仅加载参数
model.load_state_dict(torch.load(‘net_params.pkl’)) # 仅加载参数

Guardar y cargar Checkpoint para inferencia / continuar entrenando

torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)

carga

model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()
# - or -
model.train()

Conversión de modelo

libtorch no depende de python, el modelo entrenado por python debe convertirse al modelo de script antes de que libtorch pueda cargarlo

import torch
import architecture as arch
 
# An instance of your model.
model = arch.RRDB_Net(3, 3, 64, 23, gc=32, upscale=4, norm_type=None, act_type='leakyrelu', \
                        mode='CNA', res_scale=1, upsample_mode='upconv')
 
model.load_state_dict(torch.load('./models/RRDB_ESRGAN_x4.pth'), strict=True)
model.eval()
 
# An example input you would normally provide to your model's forward() method.
example = torch.rand(64, 3, 3, 3)
 
# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)
output = traced_script_module(torch.ones(64, 3, 3, 3))
traced_script_module.save("./models/RRDB_ESRGAN_x4_000.pt")
 
# The traced ScriptModule can now be evaluated identically to a regular PyTorch module
print(output)

Modelo propio

import torch
import torchvision
from PIL import Image
import numpy as np
from model import AlexNet

# 图片发在了build文件夹下
image = Image.open("/home/zhongsy/datasets/dataset/train/no_obstacle/0.jpg")
image = image.resize((224, 224), Image.ANTIALIAS)
image = np.asarray(image)
image = image / 255
image = torch.Tensor(image).unsqueeze_(dim=0)
image = image.permute((0, 3, 1, 2)).float()

model = torch.load('./AlexNet.pt',map_location=torch.device('cpu'))
model.eval()
input_cpu_ = image.cpu()
input_gpu = image.cuda()


torchd_cpu = torch.jit.trace(model, input_cpu_)
torch.jit.save(torchd_cpu, "cpu.pth")

model_gpu = torch.load('./AlexNet.pt')
model_gpu.cuda()
model_gpu.eval()

torchd_gpu = torch.jit.trace(model_gpu,input_gpu)
torch.jit.save(torchd_gpu, "gpu.pth")

Cabe señalar que si desea entrenar el modelo en la gpu y hacer inferencias en la cpu , debe transformar antes de guardar el modelo, y luego recordar llamar a model.eval (), que se ve así

• Es decir, lo que guarda la cpu solo puede ser inferencia en la cpu, es decir, lo que guarda la gpu solo puede ser inferencia en la gpu,

control de modelo onnx

import onnx

model = onnx.load('gpu.onnx')
onnx.checker.check_model(model)
print("====> pass")

Instalar libtorch

Enlace de descarga

Escribe un programa para cargar el modelo

#include <time.h>
#include <torch/script.h>

#include <algorithm>
#include <iostream>
#include <memory>
#include <opencv2/opencv.hpp>

int main() {
  clock_t start, end;
  torch::jit::script::Module model = torch::jit::load("../cpu.pth");
//   model.to(at::kCUDA);
  start = clock();

  cv::Mat input_image = cv::imread("../294.jpg");
  cv::resize(input_image, input_image, cv::Size(224, 224));

  torch::Tensor image_tensor = torch::from_blob(
      input_image.data, {input_image.rows, input_image.cols, 3}, torch::kByte);

  image_tensor = image_tensor.permute({2, 0, 1});

  image_tensor = image_tensor.toType(torch::kFloat);

  image_tensor = image_tensor.div(255);

  image_tensor = image_tensor.unsqueeze(0);

//   image_tensor = image_tensor.to(at::kCUDA);

  torch::Tensor pred = model.forward({image_tensor}).toTensor();
  end = clock();
  std::cout << "F2运行时间" << (double)(end - start) / CLOCKS_PER_SEC
            << std::endl;

  std::cout << pred << std::endl;
}

CMakeLists.txt

cmake_minimum_required(VERSION 3.4 FATAL_ERROR)
project(simnet)

set(Torch_DIR /home/zhongsy/Downloads/libtorch/share/cmake/Torch)
find_package(Torch REQUIRED)        # 查找libtorch
find_package(OpenCV REQUIRED)       # 查找OpenCV
set(CMAKE_CXX_STANDARD 14)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
if(NOT Torch_FOUND)
    message(FATAL_ERROR "Pytorch Not Found!")
endif(NOT Torch_FOUND)

message(STATUS "Pytorch status:")
message(STATUS "    libraries: ${TORCH_LIBRARIES}")

message(STATUS "OpenCV library status:")
message(STATUS "    version: ${OpenCV_VERSION}")
message(STATUS "    libraries: ${OpenCV_LIBS}")
message(STATUS "    include path: ${OpenCV_INCLUDE_DIRS}")

add_executable(simnet pytorch.cc)
target_link_libraries(simnet ${TORCH_LIBRARIES} ${OpenCV_LIBS}) 

Prueba (cpu y gpu)

Puede optar por poner datos en cpu o gpu

cpu运行时间0.094049 s
 0.4806 -0.5249
[ CPUFloatType{1,2} ]

cuda::is_available():1
Time used:33.31 ms
-10.0399  10.7939
[ CUDAFloatType{1,2} ]

Usa Tensorrt para acelerar

Enlace acelerado de tensión