V2EX = way to explore
V2EX 是一个关于分享和探索的地方
现在注册
已注册用户请  登录
MegEngineBot
V2EX  ›  推广

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

  •  
  •   MegEngineBot · 2023-02-07 13:55:40 +08:00 · 682 次点击
    这是一个创建于 640 天前的主题,其中的信息可能已经有所发展或是发生改变。

    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   
    

    3.3 训练准备

    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 ,进行高效部署。

    参考文献

    [1]: MegEngine 快速上手 [2]: Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist, 2010. [3]: Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278–2324, 1998.

    附录: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())
    
    

    C++ 推理

    前半部分我们完成了深度学习的训练,得到了 LeNet 训练权重文件,后面我们将使用训练权重文件导出静态图模型。并使用 C++调用模型完成实际部署的模拟。

    准备工作

    在上一章中,我们提到有四种保存模型的方法,如下表所示,为了训练方便起见,保存了 checkpoint 文件。但实际部署中我们经常使用静态图模型,所以我们首先要完成静态图导出。

    方法 优劣
    保存 /加载整个模型 任何情况都不推荐
    保存加载模型状态字典 适用于推理,不满足恢复训练要求
    保存加载检查点 适用于推理或恢复训练
    导出静态图模型 适用于推理,追求高性能部署

    到处静态图在 megengine 中有较完整的教程,请参考导出序列化模型文件( Dump )。主要分为三步:

    1. 将循环内的前向计算、反向传播和参数优化代码提取成单独的函数,如下面例子中的 train_func()
    2. 将网络所需输入作为训练函数的参数,并返回任意你需要的结果(如输出结果、损失函数值等);
    3. 用 jit 模块中的 trace 装饰器来装饰这个函数,将其中的代码变为静态图代码。

    在上一章最后的附录 train.py 中有 dump 静态图的方法,代码如下:

    
    from megengine import jit
    
    def dump_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)
            pred_normalized = F.softmax(pred)
            return pred_normalized
        
        input = megengine.Tensor(np.random.randn(1, 1, 32, 32))
        output = infer_func(input, model=model)
        infer_func.dump("./lenet.mge", arg_names=["input"])
    

    调用 dump_mge 方法即可完成静态图导出。

    inference 代码

    代码的主逻辑为:

    1. 创建 Network
    2. 使用 load_model()载入模型
    3. 使用 stb 预处理图片(加载和 resize),然后归一化,载入进 input tensor
    4. 使用 network->forward()和 network->wait()完成推理逻辑。
    5. 获取模型输出 tensor ,并对其进行处理。

    推理代码为:

    //inference.cpp
    #include <iostream>
    #include <stdlib.h>
    #define STB_IMAGE_IMPLEMENTATION
    #include "stb/stb_image.h"
    #define STB_IMAGE_WRITE_IMPLEMENTATION
    #include "stb/stb_image_write.h"
    #define STB_IMAGE_RESIZE_IMPLEMENTATION
    #define STB_IMAGE_RESIZE_STATIC
    #include "stb/stb_image_resize.h"
    
    #include "lite/network.h"
    #include "lite/tensor.h"
    
    //注意在这里修改测试图片与所用模型
    #define IMAGE_PATH "./test.png"
    #define MODEL_PATH "./lenet.mge"
    
    void preprocess_image(std::string pic_path, std::shared_ptr<lite::Tensor> tensor) {
        int width, height, channel;
        uint8_t* image = stbi_load(pic_path.c_str(), &width, &height, &channel, 0);
        printf("Input image %s with height=%d, width=%d, channel=%d\n", pic_path.c_str(),
               width, height, channel);
    
        auto layout = tensor->get_layout();
        auto pixels = layout.shapes[2] * layout.shapes[3];
        size_t image_size = width * height * channel;
        size_t gray_image_size = width * height * 1;
        unsigned char *gray_image = (unsigned char *)malloc(gray_image_size);
        for(unsigned char *p=image, *pg=gray_image; p!=image+image_size; p+=channel,pg++)
        {
            *pg = uint8_t(*p + *(p+1) + *(p+2))/3.0;
        }
        
        //! resize to tensor shape
        std::shared_ptr<std::vector<uint8_t>> resize_int8 =
                std::make_shared<std::vector<uint8_t>>(pixels * 1);
        stbir_resize_uint8(
                gray_image, width, height, 0, resize_int8->data(), layout.shapes[2],
                layout.shapes[3], 0, 1);
    
        free(gray_image);
        stbi_image_free(image);
        //! 减去均值,归一化
        unsigned int sum = 0;
        for(unsigned char *p=gray_image; p!=gray_image+gray_image_size;p++){
        sum += *p;
        }
        sum /= gray_image_size;
        float* in_data = static_cast<float*>(tensor->get_memory_ptr());
        for (size_t i = 0; i < pixels; i++) {
            in_data[i] = resize_int8->at(i)-sum;     
        }
    }
    
    int main()
    {
        //创建网络
        std::shared_ptr<lite::Network> network = std::make_shared<lite::Network>();
        //加载模型
        network->load_model(MODEL_PATH);
        
        std::shared_ptr<lite::Tensor> input_tensor = network->get_io_tensor("input");
        preprocess_image(IMAGE_PATH, input_tensor);
        //将图片转为 Tensor
        
        network->forward();
        network->wait();
    
        std::shared_ptr<lite::Tensor> output_tensor = network->get_output_tensor(0);
        float* predict_ptr = static_cast<float*>(output_tensor->get_memory_ptr());
        float max_prob = predict_ptr[0];
        size_t number = 0;
        //寻找最大的标签
        for(size_t i=0; i<10; i++)
        {
            float cur_prob = predict_ptr[i];
            if(cur_prob>max_prob)
            {
                max_prob = cur_prob;
                number = i;
            }
        }
        std::cout << "the predict number is :" << number << std::endl;
        return 0;
    }
    

    推理的代码已经编写完成,还需要对其进行编译,根据我们部署的平台,选择编译方式,比如安卓,可以选择交叉编译。这里我们选择部署在本机上。

    可以使用 g++进行编译,编译时需要连接 MegEngine Lite 库文件,并且准备好 stb 头文件。

    配置环境

    由于使用 C++调用 MegEngine Lite 接口,所以我们首先需要编译出 MegEngine Lite 的库。 安装 MegEngine:从源代码编译 MegEngine 。请参考编译 MegEngine Lite

    1. clone MegEngine 工程,进入根目录
    git clone --depth=1 [email protected]:MegEngine/MegEngine.git
    cd MegEngine
    
    1. 安装 MegEngine 所需的依赖
    ./third_party/prepare.sh
    ./third_party/install-mkl.sh
    
    1. 使用 cmake 进行编译工程得到 c++推理所需的库文件
    scripts/cmake-build/host_build.sh
    

    编译完成后,需要的库文件所在地址为:

    MegEngine/build_dir/host/MGE_WITH_CUDA_OFF/MGE_INFERENCE_ONLY_ON/Release/install/lite/

    这里为了在 g++编译时添加库文件方便,可以将库文件地址设为环境变量

    export LITE_INSTALL_DIR=/path/to/megenginelite-lib #上一步中编译生成的库文件安装路径
    export LD_LIBRARY_PATH=$LITE_INSTALL_DIR/lib/x86_64/:$LD_LIBRARY_PATH
    

    安装 stb: stb 是一个轻量化的图片加载库,可以替代 opencv 完成图片的解码。想要使用它,只需要将对应的头文件包含到项目内,不像 opencv 需要编译产生链接库。

    这里为了调用方便直接将 stb 的项目下载下来:

    git clone https://github.com/nothings/stb.git
    

    想要使用图片加载函数 stbi_load(),只需在 cpp 文件中 define STB_IMAGE_IMPLEMENTATION 并且 include stb_image.h 头文件

    #define STB_IMAGE_IMPLEMENTATION
    #include "stb/stb_image.h"
    

    动态链接编译

    最后使用 g++或者 clang 完成对 inference.cpp 的编译。

    • -I 选项添加编译时头文件搜索路径
    • -l 添加动态链接库
    • -L 添加动态链接库搜索路径
    g++ -o inference -I$LITE_INSTALL_DIR/include -I./stb inference.cpp -llite_shared -L$LITE_INSTALL_DIR/lib/x86_64
    

    编译后会在本目录下会得到 inference 二进制文件

    执行二进制文件

    准备好一张手写数字图片,将图片与模型放到同一目录,执行编译好的文件即可得到推理结果。

    ./inference test.jpg
    

    以上就完成了 LeNet 神经网络的部署。

    更多 MegEngine 信息获取,您可以: 查看 MegEngine 官网 GitHub 项目,或加入 MegEngine 用户交流 QQ 群:1029741705

    目前尚无回复
    关于   ·   帮助文档   ·   博客   ·   API   ·   FAQ   ·   实用小工具   ·   3268 人在线   最高记录 6679   ·     Select Language
    创意工作者们的社区
    World is powered by solitude
    VERSION: 3.9.8.5 · 26ms · UTC 10:42 · PVG 18:42 · LAX 02:42 · JFK 05:42
    Developed with CodeLauncher
    ♥ Do have faith in what you're doing.