乐趣区

关于深度学习:深度学习快速上手基于-MegEngine-的-LeNet-快速训练与部署

jiezi

|本文转载自知乎  @嘿呀嘿   集体 blog。

1. 前言

Megengine 是旷视科技开发的一款训练推理一体化的深度学习框架，相似于 pytorch，tensorflow。

应用 Megengine 能够疾速实现常见的深度学习模型，本文将应用 Megengine 实现手写数字辨认，以实现深度学习的两大步骤：训练和预测。通过本文，读者对深度学习的最根本流程和 Megengine 框架的应用办法有大抵理解。

2. 环境装置

在命令行输出下列语句即可装置 Megengine，倡议应用 python 版本为 3.5 到 3.8

python3 -m pip install --upgrade pip
python3 -m pip install megengine -f https://megengine.org.cn/whl/mge.html

装置实现后能够在命令行测试是否装置胜利。

python3
import megengine
print(megengine.__version__)

3. 训练

本局部训练代码来自 Megengine 官网教程，须要具体理解细节请返回 MegEngine 疾速上手

3.1 数据集筹备

3.1.1 下载数据集

深度学习的第一步为筹备数据集，通常会为数据集写一个接口来拜访数据集，并对数据进行预处理。

Megengine 中曾经实现了 MNIST 数据集的接口，咱们能够通过以下代码间接获取。如果想要制作或应用其余数据集，能够点击这里进行学习。

from megengine.data.dataset import MNIST
DATA_PATH = "./datasets/MNIST"
#第一次运行后，将 download 改为 False
train_dataset = MNIST(DATA_PATH, train=True, download=True)
test_dataset = MNIST(DATA_PATH, train=False, download=True)

3.1.2 数据加载及预处理

下面应用 MNIST（）实现数据集的加载和 Dataset 的构建，接下来将对数据进行加载，批改数据。应用 DataLoader、Sampler 和 Transform 实现。

DataLoader

性能 ：构建可迭代的数据装载器，非常灵活地从数据集间断获取小批量数据
参数：

  • dataset – 须要从中分批加载的数据集。
  • sampler (Optional) – 定义从数据集中采样数据的策略。
  • transform (Optional) – 定义抽样批次的转换策略。对数据须要作的变换 默认：None
  • …
RandomSampler

性能 ：创立一个列表，蕴含所有数据的索引，可实现数据的随机取样
参数：

  • dataset – 待采样的指标数据集。
  • batch_size – 应用 batch 办法时指定 batch 大小。
  • drop_last – 如果 batch 大小不能整除数据集大小时，为 True 则放弃最初一个不残缺的 batch; 为 False 则最初一个 batch 可能比拟小。默认：False
  • …
import megengine.data as data
import megengine.data.transform as T
train_sampler = data.RandomSampler(train_dataset, batch_size=64)
test_sampler = data.SequentialSampler(test_dataset, batch_size=4)

transform = T.Compose([T.Normalize(0.1307*255, 0.3081*255),
    T.Pad(2),
    T.ToMode("CHW"),
])

train_dataloader = data.DataLoader(train_dataset, train_sampler, transform)
test_dataloader = data.DataLoader(test_dataset, test_sampler, transform)

3.2 模型

接下来定义网络结构，LeNet 的网络结构如下图所示。
[图片上传失败 …(image-41aaca-1660811354056)]
定义网络结构次要为两步：定义网络子模块和连贯网络子模块。如下代码所示，应用 init 办法创立子模块，forward()办法连接子模块。

import megengine.functional as F
import megengine.module as M
class LeNet(M.Module):
    def __init__(self):
        super().__init__()
        #输出大小为(batch, 1, 32, 32), 输入大小为(batch, 6, 28, 28)
        self.conv1 = M.Conv2d(1, 6, 5)
        self.conv2 = M.Conv2d(6, 16, 5)
        self.fc1   = M.Linear(16*5*5, 120)
        self.fc2   = M.Linear(120, 84)
        self.fc3   = M.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = F.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
  • class LeNet 继承自 Module 的类, 想要理解更多对于构建模型的细节，能够参考参考 module 定义模型构造

3.3 训练筹备

  • GradManager 负责计算梯度，治理须要优化的参数，用 attach()办法增加，退出的参数才会被计算保留梯度值。对梯度计算理解更多请点击 Autodiff 基本原理与应用
  • optimizer 抉择应用的优化办法，这里应用 Optimizer.SGD（随机梯度降落法）作为优化器，对优化器理解更多请点击应用 optimizer 优化参数
import megengine.optimizer as optim
import megengine.autodiff as autodiff
gm = autodiff.GradManager().attach(model.parameters())
#参数为须要优化的参数，学习率等
optimizer = optim.SGD(model.parameters(),
    lr=0.01,
    momentum=0.9,
    weight_decay=5e-4
)

3.4 训练迭代

接下来进入程序的主逻辑，开始训练模型。应用两个嵌套循环，一个大循环为一个 epoch，遍历一次数据集，计算一次准确度。

每个小循环为一个 batch，将一批数据传入模型中，进行前向计算失去预测概率，应用穿插熵 (cross_entropy) 来计算 loss, 接着调用 GradManager.backward 办法进行反向计算并且记录每个 tensor 的梯度信息。而后应用 Optimizer.step 办法更新模型中的参数。因为每次更新参数后不主动革除梯度，所以还须要调用 clear_grad 办法。

import megengine
epochs = 10
model.train()
for epoch in range(epochs):
    total_loss = 0
    for batch_data, batch_label in train_dataloader:
        batch_data = megengine.Tensor(batch_data)
        batch_label = megengine.Tensor(batch_label)

        with gm:
            logits = model(batch_data)
            loss = F.nn.cross_entropy(logits, batch_label)
            gm.backward(loss)
            optimizer.step().clear_grad()

        total_loss += loss.item()

    print(f"Epoch: {epoch}, loss: {total_loss/len(train_dataset)}")

3.5 保留模型

罕用的神经网络都具备十分大数量级的参数，每次训练须要破费很长时间，为了可能训练中断后可能依照上次训练的成绩接着训练，咱们能够每 10 个 epoch 保留一次模型（或更多）。保留模型有几种办法，如表所示。办法具体介绍请点击保留与加载模型。

办法 优劣
保留 / 加载整个模型 任何状况都不举荐
保留加载模型状态字典 实用于推理，不满足复原训练要求
保留加载检查点 实用于推理或复原训练
导出动态图模型 实用于推理，谋求高性能部署

咱们抉择保留加载检查点，既能够用于复原训练也能够推理。保留时调用 megengine.save()办法，参数如下：

megengine.save({
                "epoch": epoch,
                "state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "loss": loss,
                ...
               }, PATH)

而后就能够欢快的进行训练了，察看训练后果，当 loss 降落到肯定境地，准确率满足要求后，终止训练.

如果训练产生中断，能够调用 load()办法和 optimizer.load_state_dict()办法，对模型的加载，从新开始训练。代码如下：

model = LeNet()
optimizer = optim.SGD()

checkpoint = megengine.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()

4. 推理

下面几个章节曾经实现深度学习大部分内容，曾经可能产生一个须要的算法模型。这个算法对筹备好的数据集有比拟好的拟合成果，然而咱们的最终目标是用模型进行推理，即可能对新的数据进行预测。这将是上面介绍的内容。

首先有一种很简略的办法，应用 python 加载模型并设定 model.eval()，代码如下所示，这样就能够简略调用训练好的模型用以理论。

from train import LeNet
import cv2
import numpy as np
import megengine
import megengine.data.transform as T
import megengine.functional as F

IMAGE_PATH = "./test.png"
CHECK_POINT_PATH = "./checkpoint.pkl"

def load_model(check_point_path = CHECK_POINT_PATH):
    model  = LeNet()
    check_point = megengine.load(check_point_path)
    #留神 checkpoint 保留时模型对应的键，此处为 state_dict
    model.load_state_dict(check_point["state_dict"])
    model.eval()
    return model

def main():
    # 加载一张图像为灰度图
    image = cv2.imread(IMAGE_PATH,cv2.IMREAD_GRAYSCALE)
    image = cv2.resize(image, (32, 32))
    #将图片变换为黑底白字
    image = np.array(255-image)
    tensor_image = megengine.tensor(image).reshape(1, 1, 32, 32)
    model = load_model()
    logit= model(tensor_image)
    pred = F.argmax(logit, axis=1).item()
    print("number:", pred)

if __name__ == "__main__":
    main()

不过在理论部署中，还须要思考部署环境，推理速度等因素，所以从训练好模型到部署落地还有很长的路。Megengine 因为其设计特点——训练推理一体化，能够不便地将训练模型部署。这将是下一章介绍的内容，下一章将应用 C++ 调用 Megengine lite，进行高效部署。

参考文献

附录：train.py

from megengine.data.dataset import MNIST
from megengine import jit, tensor
import megengine
import numpy as np
import megengine.data as data
import megengine.data.transform as T
import megengine.functional as F
import megengine.module as M
import megengine.optimizer as optim
import megengine.autodiff as autodiff

DATA_PATH = "./datasets/train/"

def load_data(data_path =DATA_PATH):
    train_dataset = MNIST(DATA_PATH)
    test_dataset = MNIST(DATA_PATH)

    train_sampler = data.RandomSampler(train_dataset, batch_size=64)
    test_sampler = data.SequentialSampler(test_dataset, batch_size=2)


    transform = T.Compose([T.Normalize(0.1307*255, 0.3081*255),
        T.Pad(2),
        T.ToMode("CHW"),
    ])

    train_dataloader = data.DataLoader(train_dataset, train_sampler, transform)
    test_dataloader = data.DataLoader(test_dataset, test_sampler, transform)
    return train_dataloader, test_dataloader

#Define model
class LeNet(M.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = M.Conv2d(1, 6, 5)
        self.conv2 = M.Conv2d(6, 16, 5)
        self.fc1   = M.Linear(16*5*5, 120)
        self.fc2   = M.Linear(120, 84)
        self.fc3   = M.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = F.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


def train(dataloader):
    model = LeNet()

    # GradManager and Optimizer setting
    gm = autodiff.GradManager().attach(model.parameters())
    optimizer = optim.SGD(model.parameters(),
        lr=0.01,
        momentum=0.9,
        weight_decay=5e-4
    )


    # Training and validation
    nums_epoch = 50
    for epoch in range(nums_epoch):
        training_loss = 0
        nums_train_correct, nums_train_example = 0, 0
        nums_val_correct, nums_val_example = 0, 0

        for step, (image, label) in enumerate(dataloader[0]):
            image = megengine.Tensor(image)
            label = megengine.Tensor(label)

            with gm:
                score = model(image)
                loss = F.nn.cross_entropy(score, label)
                gm.backward(loss)
                optimizer.step().clear_grad()

            training_loss += loss.item() * len(image)

            pred = F.argmax(score, axis=1)
            nums_train_correct += (pred == label).sum().item()
            nums_train_example += len(image)

        training_acc = nums_train_correct / nums_train_example
        training_loss /= nums_train_example

        for image, label in dataloader[1]:
            image = megengine.Tensor(image)
            label = megengine.Tensor(label)
            pred = F.argmax(model(image), axis=1)

            nums_val_correct += (pred == label).sum().item()
            nums_val_example += len(image)

        val_acc = nums_val_correct / nums_val_example
        #每十次 epoch 保留一次模型
        if epoch%2 == 0:
            megengine.save(
                    {"epoch":epoch, 
                    "state_dict": model.state_dict(),
                    "optimizer_state_dict": optimizer.state_dict(),
                    "loss": loss,
                    }, 
                    "./checkpoint.pkl")
        print(f"Epoch = {epoch},"
            f"train_loss = {training_loss:.3f},"
            f"train_acc = {training_acc:.3f},"
            f"val_acc = {val_acc:.3f}")


def dumpy_mge(pkl_path = "checkpoint.pkl"):
    model = LeNet()
    check_point = megengine.load(pkl_path)
    model.load_state_dict(check_point["state_dict"])
    model.eval()

    @jit.trace(symbolic=True, capture_as_const=True)
    def infer_func(input, *, model):
        pred  = model(input)
        return pred
    
    input = megengine.Tensor(np.random.randn(1, 1, 32, 32))
    output = infer_func(input, model=model)
    infer_func.dump("./lenet.mge", arg_names=["input"])

if __name__=='__main__':
    train(load_data())

更多 MegEngine 信息获取，您能够：
查看 MegEngine 官网和 GitHub 我的项目，或退出 MegEngine 用户交换 QQ 群：1029741705

退出移动版