1. 导出模型
这里是我随便复现的一个模型 好像是vgg16 记不清了(*^▽^*)
from torch import nn
import torch
class My_AFIRModule(nn.Module):
def __init__(self):
super(My_AFIRModule, self).__init__()
self.model1 = nn.Sequential(
nn.Conv2d(3,32,5,padding=2),
nn.MaxPool2d(2),
nn.Conv2d(32,32,5,padding=2),
nn.MaxPool2d(2),
nn.Conv2d(32,64,5,padding=2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(1024,64),
nn.Linear(64,10)
)
def forward(self, x):
x = self.model1(x)
return x
if __name__ == "__main__":
net = My_AFIRModule()
input = torch.ones(64,3,32,32)
y = net(input)
print(y.shape)
pytorch已经给我们提供了保存的高级api接口,直接使用导出为onnx格式的模型
import cv2
import onnxruntime
import torchvision
import numpy as np
import torch
from PIL import Image
image = Image.open('data/dog.png')
# 2. Preprocess the image
transform = torchvision.transforms.Compose([
torchvision.transforms.Resize((32, 32)),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
image = transform(image)
image = torch.reshape(image, (1, 3, 32, 32))
ort_session = onnxruntime.InferenceSession("mymodule.onnx")
ort_inputs = {'input': image.numpy()}
print(image.numpy())
ort_output = ort_session.run(['output'], ort_inputs)[0]
ort_output = torch.from_numpy(ort_output)
_, predict = torch.max(ort_output.data, 1)
print('Inference result: ', predict.item())
2. 通过依赖库调用模型(平台随意,这里以idea为例)
推理脚本写的比较一般ε=(´ο`*))) 其中比较注意的是模型的路径,因为我们要使用Docker部署,所以是动态加载的路径,如果要本地跑起来,需要在运行时通过命令行指定资源路径resource.path
注意:在java中设置路径前加上-D
package com.example.testtwo;
import ai.onnxruntime.OnnxTensor;
import ai.onnxruntime.OrtEnvironment;
import ai.onnxruntime.OrtSession;
import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.nio.FloatBuffer;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.HashMap;
import java.util.Map;
public class OnnxModelInference {
public static void main(String[] args) {
Inference();
}
public static String Inference(){
try {
// 1. 加载图片
String resourcePath = System.getProperty("resource.path", "/opt/drag-master/resource");
BufferedImage image = ImageIO.read(new File(resourcePath + "/dog.png"));
// 2. 预处理图片
// 假设你已经有方法将BufferedImage转换为预处理后的浮点数组
float[] inputData = preprocessImage(image);
// 定义输入的形状 (1, 3, 32, 32)
long[] inputShape = new long[]{1, 3, 32, 32};
// 3. 初始化ONNX Runtime环境
OrtEnvironment env = OrtEnvironment.getEnvironment();
// 4. 加载ONNX模型
OrtSession.SessionOptions options = new OrtSession.SessionOptions();
OrtSession session = env.createSession(resourcePath + "/mymodule.onnx", options);
// 5. 准备输入数据
OnnxTensor inputTensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(inputData), inputShape);
// 6. 创建输入映射
Map<String, OnnxTensor> inputs = new HashMap<>();
inputs.put("input", inputTensor);
// 7. 执行推理
OrtSession.Result result = session.run(inputs);
// 8. 获取并处理输出
float[][] outputData = (float[][]) result.get(0).getValue(); // 假设输出是二维数组
int predictedClass = argMax(outputData[0]);
System.out.println("Inference result: " + predictedClass);
// 9. 释放资源
inputTensor.close();
session.close();
env.close();
return "Inference result: " + predictedClass;
} catch (Exception e) {
e.printStackTrace();
return "null";
}
}
// 预处理方法
private static float[] preprocessImage(BufferedImage image) {
// 使用Java来实现类似Python中torchvision.transforms的功能
int height = 32;
int width = 32;
float meanR = 0.485f, meanG = 0.456f, meanB = 0.406f;
float stdR = 0.229f, stdG = 0.224f, stdB = 0.225f;
BufferedImage resizedImage = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
resizedImage.getGraphics().drawImage(image, 0, 0, width, height, null);
float[] data = new float[3 * width * height];
for (int y = 0; y < height; y++) {
for (int x = 0; x < width; x++) {
int rgb = resizedImage.getRGB(x, y);
int r = (rgb >> 16) & 0xFF;
int g = (rgb >> 8) & 0xFF;
int b = rgb & 0xFF;
data[y * width + x] = ((r / 255.0f) - meanR) / stdR;
data[height * width + y * width + x] = ((g / 255.0f) - meanG) / stdG;
data[2 * height * width + y * width + x] = ((b / 255.0f) - meanB) / stdB;
}
}
return data;
}
// 寻找最大值索引的方法
private static int argMax(float[] array) {
int maxIndex = 0;
float maxValue = array[0];
for (int i = 1; i < array.length; i++) {
if (array[i] > maxValue) {
maxValue = array[i];
maxIndex = i;
}
}
return maxIndex;
}
}
2.1. 新建Spring项目(合并脚本)
新建ApiController类 创建一个get接口
package com.example.testtwo;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import java.io.IOException;
@RestController
public class ApiController {
@GetMapping ("/infer")
public String infer() throws IOException {
// Post-process the output
String result = OnnxModelInference.Inference();
return result;
}
}
2.2 将项目导出为jar包
因为是maven项目,所以在右侧的maven选项中依次点击
导出的包位于当前项目的target目录下,将其复制到一个新的干净的目录下,我们来构建Docker的镜像文件。
3. 构建Docker镜像
在新的目录下新建Dockerfile文件,没有扩展名!没有扩展名!而且文件名不要改,不要改!
FROM openjdk:17
ENV TZ=Asia/Shanghai
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
WORKDIR /opt/drag-master
COPY Inference.jar /opt/drag-master/
EXPOSE 8080
CMD ["java", "-Dresource.path=/opt/drag-master/resource","-jar", "Inference.jar"]
首先是java环境的导入,使用了openjdk的镜像,可以根据自己的项目进行版本修改,
然后是设置时区,这里不需要修改,(对应第二、第三行)
新建工作目录
复制jar包到工作目录,这里根据自己的jar包名进行修改
端口设置,随意,不要使用常用的端口就行,防止端口冲突
CMD这里我们新建了resource.path的环境变量,原本我把模型文件直接copy进去,发现调用资源一直失败,所以采用的数据卷挂载的方式。
dockerfile文件写好后保存退出。
我是windows下完成的这些操作 ,提一嘴因为众所周知的原因,在构建镜像时可能会出现问题(我的是使用了魔法工具)
在
在当前目录下打开cmd终端,window下安装Docker桌面软件一定要运行该软件才行(只有这样cmd才能运行Docker指令)
-t 后面是构建的镜像的名称 . 一定不能省略,这是告诉其在当前目录查找dockerfile文件
docker build -t yourpagename .
然后通过以下命令查看构建的镜像
docker images
根据构建的镜像名称进行镜像的导出/保存
-o 后是我们导出后的名称,最后是要保存的镜像名 :后是版本号(不用修改)
docker save -o outputimagename imagename:latest
其输出位置在当前终端的路径。
至此docker镜像构建完成,我在华为1核2G的服务器部署了一下,还可以。
先将导出的镜像,上传到服务器上,通过下列命令进行加载
docker load -i iamgepath
然后通过之前的docker images 命令进行查看
尝试了好多次o(╥﹏╥)o
然后建立容器并启动,运行下列的命令即可
这里我将导出的模型以及要使用的照片放在了服务器的/root/resource路径下,并且文件名要和java的推理脚本中使用的文件名保持一致(这里只是一个简单的例子,所以我写死了),通过数据卷的方式进行资源的访问。命令的最后一块demoone 将其修改成你自己的镜像名称即可
docker run -d --name demotar -v /root/resource:/opt/drag-master/resource -p 8080:8080 demoone
通过docker ps 可以查看启动的容器状态
通过ip地址+端口号+接口进行访问
部署的是一个分类模型,输出结果为类别对应的顺序代号,嗯,(结果好像是错的,就跑了几轮。。。)