使用Docker部署模型示例

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地址+端口号+接口进行访问

部署的是一个分类模型，输出结果为类别对应的顺序代号，嗯，（结果好像是错的，就跑了几轮。。。）