｜本文转载自知乎 @嘿呀嘿 集体 blog。
1. 前言
Megengine 是旷视科技开发的一款训练推理一体化的深度学习框架，相似于 pytorch，tensorflow。
应用 Megengine 能够疾速实现常见的深度学习模型，本文将应用 Megengine 实现手写数字辨认，以实现深度学习的两大步骤：训练和预测。通过本文，读者对深度学习的最根本流程和 Megengine 框架的应用办法有大抵理解。
2. 环境装置
在命令行输出下列语句即可装置 Megengine ，倡议应用 python 版本为 3.5 到 3.8
python3 -m pip install --upgrade pippython3 -m pip install megengine -f https://megengine.org.cn/whl/mge.html装置实现后能够在命令行测试是否装置胜利。
python3import megengineprint(megengine.__version__)3. 训练
本局部训练代码来自 Megengine 官网教程，须要具体理解细节请返回 MegEngine 疾速上手
3.1 数据集筹备
3.1.1 下载数据集
深度学习的第一步为筹备数据集，通常会为数据集写一个接口来拜访数据集，并对数据进行预处理。
Megengine 中曾经实现了 MNIST 数据集的接口，咱们能够通过以下代码间接获取。如果想要制作或应用其余数据集，能够点击这里进行学习。
from megengine.data.dataset import MNISTDATA_PATH = "./datasets/MNIST"#第一次运行后，将 download 改为 Falsetrain_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 dataimport megengine.data.transform as Ttrain_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 Fimport megengine.module as Mclass 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 optimimport megengine.autodiff as autodiffgm = 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 megengineepochs = 10model.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 LeNetimport cv2import numpy as npimport megengineimport megengine.data.transform as Timport megengine.functional as FIMAGE_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 modeldef 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 MNISTfrom megengine import jit, tensorimport megengineimport numpy as npimport megengine.data as dataimport megengine.data.transform as Timport megengine.functional as Fimport megengine.module as Mimport megengine.optimizer as optimimport megengine.autodiff as autodiffDATA_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 modelclass 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 xdef 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