本文共 817 字，大约阅读时间需要 2 分钟。

Libtorch常用代码

PyTorch生成模型

预处理部分

张量的常见操作

创建张量

auto t = torch::zeros({3,4});

t = torch::ones({3,4});

t = torch::tensor({33,22,11});

张量运算

模型加载与预测

加载预训练模型

from torchvision.models import resnet34

import torch.nn.functional as F

import torch.nn as nn

import torch

import cv2

读取图片并归一化

image = cv2.imread("flower.jpg")

image = cv2.resize(image, (224, 224))

input_tensor = torch.tensor(image).permute(2, 0, 1).unsqueeze(0).float() / 225.0

加载模型

model = resnet34(pretrained=True)

模型评估模式

model.eval()

进行预测

output = model(input_tensor)

output = F.softmax(output, 1)

展示结果

print("模型预测结果为第{}类，置信度为{}".format(torch.argmax(output), output.max()))

模型生成(.pt文件)

model = model.to(torch.device("cpu"))

model.eval()

var = torch.ones((1, 3, 224, 224))

traced_script_module = torch.jit.trace(model, var)

traced_script_module.save("resnet34.pt")

转载地址：http://owwfk.baihongyu.com/