Pytorch项目打包和部署（2）——代码——flask+socket+pt2pt

一、flask

1.1 module.py

import torch
from torch import nn

class TestModel(nn.Module):

    def __init__(self):
        super(TestModel, self).__init__()
        self.sequential = nn.Sequential(
            nn.Linear(784, 100),
            nn.LeakyReLU(0.1),
            nn.Linear(100, 10)
        )

    # @torch.jit.script_method
    def forward(self, x):
        return self.sequential(x)

if __name__ == '__main__':
    input = torch.randn(1, 784)
    model = TestModel()

    model.eval()
    traced_script_module = torch.jit.trace(model, input)
    traced_script_module.save("model.pt")

1.2 flash_test.py

from flask import Flask
import json
import torch
from flask_Demo import module

app = Flask(__name__)#确保Flask调用的是当前模块
net = module.TestModel()
#使用网络
net.eval()
# print(torch.cuda.get_device_name())
"""
在python中以@开头的函数称为装饰器，装饰器就是接收函数的函数,
装饰器和普通函数不一样，普通函数接收的是变量，返回的是一个或多个值，
而装饰器接收的是一个函数，返回的也是一个函数，而非函数的结果。
"""
"""
普通装饰器不能接收任何参数，但是@app.route()是需要接收参数的
在代码运行期间，动态的增加函数的功能的方式，称为“装饰器”。
"""
@app.route("/")
def hello():
    x = torch.randn(2,784)
    y = net(x)
    _y = y.detach().numpy()
    return json.dumps(_y.tolist())#将字典、列表转化为字符串

if __name__ == '__main__':
    app.run()

1.3 MTCNN侦测网络的flask通信流程

# MTCNN的使用
# 流程：图像-->缩放-->P网(NMS和边界框回归)-->R网路(NMS和边界框回归)--->O网络(NMS和边界框回归)

import torch
from PIL import Image
from PIL import ImageDraw
import numpy as np
from tool import utils
import nets
from torchvision import transforms
import time
import os


# 网络调参
# P网络:
p_cls = 0.6 #原为0.6
p_nms = 0.5 #原为0.5
# R网络：
r_cls = 0.6 #原为0.6
r_nms = 0.5 #原为0.5
# R网络：
o_cls = 0.3 #原为0.97
o_nms = 0.5 #原为0.7


# 侦测器
class Detector:
    # 初始化时加载三个网络的权重(训练好的)，cuda默认设为True
    def __init__(self, pnet_param="./param_0/pnet.pt", rnet_param="./param_0/rnet.pt", onet_param="./param_0/onet.pt",
                 isCuda=True):

        self.isCuda = isCuda

        self.pnet = nets.PNet() # 创建实例变量，实例化P网络
        self.rnet = nets.RNet()
        self.onet = nets.ONet()

        if self.isCuda:
            self.pnet.cuda() # 给P网络加速
            self.rnet.cuda()
            self.onet.cuda()

        self.pnet.load_state_dict(torch.load(pnet_param)) # 把训练好的权重加载到P网络中
        self.rnet.load_state_dict(torch.load(rnet_param))
        self.onet.load_state_dict(torch.load(onet_param))

        self.pnet.eval() # 训练网络里有BN（批归一化时），要调用eval方法，使用是不用BN，dropout方法
        self.rnet.eval()
        self.onet.eval()

        self.__image_transform = transforms.Compose([
            transforms.ToTensor()
        ]) # 图片数据类型转换

    def detect(self, image): # 检测图片

        # P网络检测-----1st
        start_time = time.time()               # 开始计时
        pnet_boxes = self.__pnet_detect(image) # 调用__pnet_detect函数（后面定义）
        if pnet_boxes.shape[0] == 0:           # 若P网络没有人脸时，避免数据出错，返回一个新数组
            return np.array([])
        end_time = time.time()                 # 计时结束
        t_pnet = end_time - start_time         # P网络所占用的时间差
        # return pnet_boxes                    # p网络检测出的框

        # R网络检测-------2nd
        start_time = time.time()
        rnet_boxes = self.__rnet_detect(image, pnet_boxes) # 传入原图，P网络的一些框，根据这些框在原图上抠图
        if rnet_boxes.shape[0] == 0:
            return np.array([])
        end_time = time.time()
        t_rnet = end_time - start_time
        # return rnet_boxes

        #O网络检测--------3rd
        start_time = time.time()
        onet_boxes = self.__onet_detect(image, rnet_boxes) # 把原图和R网络里的框传到O网络里去
        if onet_boxes.shape[0] == 0:                      # 若P网络没有人脸时，避免数据出错，返回一个新数组
            return np.array([])
        end_time = time.time()
        t_onet = end_time - start_time

        # 三网络检测的总时间
        t_sum = t_pnet + t_rnet + t_onet
        print("total:{0} pnet:{1} rnet:{2} onet:{3}".format(t_sum, t_pnet, t_rnet, t_onet))

        return onet_boxes

    # 创建P网检测函数
    def __pnet_detect(self, image): # ★p网络全部是卷积，与输入图片大小无关，可输出任意形状图片
        boxes = [] # 创建空列表,接收符合条件的建议框

        img = image
        w, h = img.size
        min_side_len = min(w, h) # 获取图片的最小边长

        scale = 1 # 初始缩放比例（为1时不缩放）:得到不同分辨率的图片
        while min_side_len > 12: # 直到缩放到小于等于12时停止
            img_data = self.__image_transform(img) # 将图片数组转成张量
            if self.isCuda:
                img_data = img_data.cuda() # 将图片tensor传到cuda里加速
            img_data.unsqueeze_(0) # 在“批次”上升维（测试时传的不止一张图片）
            # print("img_data:",img_data.shape) # [1, 3, 416, 500]：C=3,W=416,H=500

            _cls, _offest = self.pnet(img_data) # ★★返回多个置信度和偏移量
            # print("_cls",_cls.shape)         # [1, 1, 203, 245]:NCWH：分组卷积的特征图的通道和尺寸★
            # print("_offest", _offest.shape) # [1, 4, 203, 245]:NCWH

            cls= _cls[0][0].cpu().data  # [203, 245]：分组卷积特征图的尺寸：W,H
            offest = _offest[0].cpu().data  #[4, 203, 245] # 分组卷积特征图的通道、尺寸:C,W,H
            idxs = torch.nonzero(torch.gt(cls, p_cls)) # ★置信度大于0.6的框索引；把P网络输出，看有没没框到的人脸，若没框到人脸，说明网络没训练好；或者置信度给高了、调低

            for idx in idxs: # 根据索引，依次添加符合条件的框；cls[idx[0], idx[1]]在置信度中取值：idx[0]行索引，idx[1]列索引
                boxes.append(self.__box(idx, offest, cls[idx[0], idx[1]], scale)) # ★调用框反算函数_box（把特征图上的框，反算到原图上去），把大于0.6的框留下来；

            scale *= 0.7 # 缩放图片:循环控制条件
            _w = int(w * scale) # 新的宽度
            _h = int(h * scale)

            img = img.resize((_w, _h)) # 根据缩放后的宽和高，对图片进行缩放
            min_side_len = min(_w, _h) # 重新获取最小宽高

        return utils.nms(np.array(boxes), p_nms) #返回框框，原阈值给p_nms=0.5（iou为0.5），尽可能保留IOU小于0.5的一些框下来，若网络训练的好，值可以给低些

    # 特征反算：将回归量还原到原图上去，根据特征图反算的到原图建议框
    def __box(self, start_index, offset, cls, scale, stride=2, side_len=12): # p网络池化步长为2

        _x1 = (start_index[1].float() * stride) / scale # 索引乘以步长，除以缩放比例；★特征反算时“行索引，索引互换”，原为[0]
        _y1 = (start_index[0].float() * stride) / scale
        _x2 = (start_index[1].float() * stride + side_len) / scale
        _y2 = (start_index[0].float() * stride + side_len) / scale

        ow = _x2 - _x1  # 人脸所在区域建议框的宽和高
        oh = _y2 - _y1

        _offset = offset[:, start_index[0], start_index[1]] # 根据idxs行索引与列索引，找到对应偏移量△δ:[x1,y1,x2,y2]
        x1 = _x1 + ow * _offset[0] # 根据偏移量算实际框的位置，x1=x1_+w*△δ；生样时为:△δ=x1-x1_/w
        y1 = _y1 + oh * _offset[1]
        x2 = _x2 + ow * _offset[2]
        y2 = _y2 + oh * _offset[3]

        return [x1, y1, x2, y2, cls]  # 正式框：返回4个坐标点和1个偏移量

    # 创建R网络检测函数
    def __rnet_detect(self, image, pnet_boxes):

        _img_dataset = [] # 创建空列表，存放抠图
        _pnet_boxes = utils.convert_to_square(pnet_boxes) # ★给p网络输出的框，找出中心点，沿着最大边长的两边扩充成“正方形”，再抠图
        for _box in _pnet_boxes: # ★遍历每个框，每个框返回框4个坐标点，抠图，放缩，数据类型转换，添加列表
            _x1 = int(_box[0])
            _y1 = int(_box[1])
            _x2 = int(_box[2])
            _y2 = int(_box[3])

            img = image.crop((_x1, _y1, _x2, _y2)) # 根据4个坐标点抠图
            img = img.resize((24, 24)) # 放缩在固尺寸
            img_data = self.__image_transform(img) # 将图片数组转成张量
            _img_dataset.append(img_data)

        img_dataset =torch.stack(_img_dataset) # stack堆叠(默认在0轴)，此处相当数据类型转换，见例子2★
        if self.isCuda:
            img_dataset = img_dataset.cuda() # 给图片数据采用cuda加速

        _cls, _offset = self.rnet(img_dataset) # ★★将24*24的图片传入网络再进行一次筛选

        cls = _cls.cpu().data.numpy() # 将gpu上的数据放到cpu上去，在转成numpy数组
        offset = _offset.cpu().data.numpy()
        # print("r_cls:",cls.shape)  # (11, 1):P网络生成了11个框
        # print("r_offset:", offset.shape)  # (11, 4)

        boxes = [] #R 网络要留下来的框，存到boxes里
        idxs, _ = np.where(cls > r_cls) # 原置信度0.6是偏低的，时候很多框并没有用(可打印出来观察)，可以适当调高些；idxs置信度框大于0.6的索引；★返回idxs:0轴上索引[0,1]，_:1轴上索引[0,0]，共同决定元素位置，见例子3
        for idx in idxs: # 根据索引，遍历符合条件的框；1轴上的索引，恰为符合条件的置信度索引（0轴上索引此处用不到）
            _box = _pnet_boxes[idx]
            _x1 = int(_box[0])
            _y1 = int(_box[1])
            _x2 = int(_box[2])
            _y2 = int(_box[3])

            ow = _x2 - _x1 # 基准框的宽
            oh = _y2 - _y1

            x1 = _x1 + ow * offset[idx][0] # 实际框的坐标点
            y1 = _y1 + oh * offset[idx][1]
            x2 = _x2 + ow * offset[idx][2]
            y2 = _y2 + oh * offset[idx][3]

            boxes.append([x1, y1, x2, y2, cls[idx][0]]) # 返回4个坐标点和置信度

        return utils.nms(np.array(boxes), r_nms) # 原r_nms为0.5（0.5要往小调），上面的0.6要往大调;小于0.5的框被保留下来

    # 创建O网络检测函数
    def __onet_detect(self, image, rnet_boxes):

        _img_dataset = [] # 创建列表，存放抠图r
        _rnet_boxes = utils.convert_to_square(rnet_boxes) # 给r网络输出的框，找出中心点，沿着最大边长的两边扩充成“正方形”
        for _box in _rnet_boxes: # 遍历R网络筛选出来的框，计算坐标，抠图，缩放，数据类型转换，添加列表，堆叠
            _x1 = int(_box[0])
            _y1 = int(_box[1])
            _x2 = int(_box[2])
            _y2 = int(_box[3])

            img = image.crop((_x1, _y1, _x2, _y2)) # 根据坐标点“抠图”
            img = img.resize((48, 48))
            img_data = self.__image_transform(img) # 将抠出的图转成张量
            _img_dataset.append(img_data)

        img_dataset = torch.stack(_img_dataset) # 堆叠，此处相当数据格式转换，见例子2
        if self.isCuda:
            img_dataset = img_dataset.cuda()

        _cls, _offset = self.onet(img_dataset)
        cls = _cls.cpu().data.numpy()       # (1, 1)
        offset = _offset.cpu().data.numpy() # (1, 4)

        boxes = [] # 存放o网络的计算结果
        idxs, _ = np.where(cls > o_cls) # 原o_cls为0.97是偏低的，最后要达到标准置信度要达到0.99999，这里可以写成0.99998，这样的话出来就全是人脸;留下置信度大于0.97的框；★返回idxs:0轴上索引[0]，_:1轴上索引[0]，共同决定元素位置，见例子3
        for idx in idxs: # 根据索引，遍历符合条件的框；1轴上的索引，恰为符合条件的置信度索引（0轴上索引此处用不到）
            _box = _rnet_boxes[idx] # 以R网络做为基准框
            _x1 = int(_box[0])
            _y1 = int(_box[1])
            _x2 = int(_box[2])
            _y2 = int(_box[3])

            ow = _x2 - _x1 # 框的基准宽，框是“方”的，ow=oh
            oh = _y2 - _y1

            x1 = _x1 + ow * offset[idx][0] # O网络最终生成的框的坐标；生样，偏移量△δ=x1-_x1/w*side_len
            y1 = _y1 + oh * offset[idx][1]
            x2 = _x2 + ow * offset[idx][2]
            y2 = _y2 + oh * offset[idx][3]

            boxes.append([x1, y1, x2, y2, cls[idx][0]]) #返回4个坐标点和1个置信度

        return utils.nms(np.array(boxes), o_nms, isMin=True) # 用最小面积的IOU；原o_nms(IOU)为小于0.7的框被保留下来
        
from flask import Flask
import json
app = Flask(__name__)
import torch

@app.route('/')
def hello():
    image_path = r"test_images"
    for i in os.listdir(image_path):
        detector = Detector()
        with Image.open(os.path.join(image_path, i)) as im:  # 打开图片
            # boxes = detector.detect(im)
            print("----------------------------")
            boxes = detector.detect(im)
            print("size:", im.size)
            imDraw = ImageDraw.Draw(im)
            for box in boxes:  # 多个框，没循环一次框一个人脸
                x1 = int(box[0])
                y1 = int(box[1])
                x2 = int(box[2])
                y2 = int(box[3])

                print((x1, y1, x2, y2))

                print("conf:", box[4])  # 置信度

        return json.dumps(boxes.tolist())

if __name__ == '__main__':
    app.run()

二、Socket

python之struct详解

2.1 server_demo.py

import socket  # 导入 socket 模块
import struct
s = socket.socket()  # 创建 socket 通信对象
host = socket.gethostname()  # 获取本地服务器主机名
port = 12345  # 设置端口（网关地址：门），一个主机可以有很多个端口
s.bind((host, port))  # 让通信对象绑定主机和端口地址

s.listen(5)  # 等待客户端连接，每隔5秒监听一次
# output = struct.pack("hello world!")

# byte=b'hello world!'
byte=bytes('hello world!',encoding='utf-8')
# byte=struct.pack("i",2)
while True:
    # 建立客户端连接，等待客户端访问，直到有客户端访问的时候，拿到访问者的地址
    c, addr = s.accept()  #
    print('连接地址：', addr)#客户端的地址
    c.send(byte)#向客户端发送信息
    c.close()  # 关闭连接

2.2 client_demo.py

import socket  # 导入 socket 模块
import struct
s = socket.socket()  # 创建 socket 通信对象
host = socket.gethostname()  # 获取本地服务器主机名
port = 12345  # 设置端口（网关地址：门），一个主机可以有很多个端口
s.bind((host, port))  # 让通信对象绑定主机和端口地址

s.listen(5)  # 等待客户端连接，每隔5秒监听一次
# output = struct.pack("hello world!")

# byte=b'hello world!'
byte=bytes('hello world!',encoding='utf-8')
# byte=struct.pack("i",2)
while True:
    # 建立客户端连接，等待客户端访问，直到有客户端访问的时候，拿到访问者的地址
    c, addr = s.accept()  #
    print('连接地址：', addr)#客户端的地址
    c.send(byte)#向客户端发送信息
    c.close()  # 关闭连接

2.3 socket实现的对话

2.3.1 server.py

#服务器端
import socket
s = socket.socket()
host = socket.gethostname()
port = 8000
s.bind((host, port))

s.listen(5)
while True:
    print("="*100)
    print('等待用户连接……')
    c, addr = s.accept() #阻塞进程，等待客户端接入
    print('用户已连接 address:{}\n{}\n服务器：请问需要什么帮助'.format(addr,"-"*50))
    c.send(bytes("请问需要什么帮助?输入quit结束对话",encoding="utf-8"))
    while True:
        mes = str(c.recv(1024),encoding = "utf-8")
        print("客户端：{}".format(mes))
        if mes == "quit":
            c.send(bytes("再见",encoding="utf-8"))
            break
        mes = input("服务器：")
        c.send(bytes(mes,encoding="utf-8"))
    c.close()

2.3.2 client.py

#客户端
import socket
s = socket.socket()
host = socket.gethostname()
port = 8000
s.connect((host, port)) #客户端接入
while True:
    mes = str(s.recv(1024),encoding = "utf-8")
    print("服务器：{}".format(mes))
    if mes == "再见":
        break
    mes = input("客户端：")
    s.send(bytes(mes,encoding="utf-8"))
s.close()

"""
若是你使用上面的代码，在不同的计算机上运行，必然是会报错的，
那是因为，我们的服务器和客户端的host都是socket.gethostname()，
这句代码是用来获取本机的hostname的，
很显然，如果你想实现两台设备之间的通讯，则必须把客户端的host改为服务器端的host。
如果某一方连续发送信息，改如何解决？
"""

三、MNIST打包

3.1 模型打包的时候注意的问题

1、如果模型中有dropout或者batchnormal的话，一定要先将模型设置为eval模式，再保存，否则在用libtorch调用后会出现随机干扰；
2、example这个张量的尺寸务必与你自己的模型的输入尺寸一直，否则会出现错误。
3、如果代码中有if条件控制，尽量避免使用torch.jit.trace来转换代码，因为它不能处理变化条件，如果非要用trace的话，可以把if条件控制改成别的形式；
4、jit不能转换第三方Python库中的函数，尽量所有代码都使用pytorch实现，如果速度不理想的话，可以参考github上的pytorch/extension-script项目，用C++实现需要的功能，然后注册成jit操作，最后转成torchscript。

import torch
import torch.jit

def pth2pt():
    device = torch.device('cpu')  # 使用cpu进行推理
    model = torch.load("./mnist_net.pth").to(device)  # 加载模型
    model.eval()  # 把模型转为test模式
    example = torch.randn([1, 1, 28, 28], dtype=torch.float32)
    traced_net = torch.jit.trace(model, example)
    traced_net.save("model.pt")
    print("模型序列化导出成功")

pth2pt()