使用model hub API
加载预训练模型API
读入和预处理图片数据
模型推理
推理预测和生成目标检测图片
下载源码,MegStudio中终端运行
MegEngine Models
综述
安装和环境配置
使用预训练模型
检测标注结果
from megengine import hub
model = hub.load(
"megengine/models", # Github相对路径
"retinanet_res50_coco_1x_800size", # 指定模型名称
pretrained=True,
use_cache=False,
)
model.eval() # 设置为评价验证模式
models_api = hub.import_module(
"megengine/models",
git_host="github.com",
)
import cv2
image = cv2.imread("./workspace/cat/cat4.jpg") # 图片的相对路径
data, im_info = models_api.DetEvaluator.process_inputs(image, 800, 1333)
# 图片裁剪为宽<=800 高<=1333
# 将处理后的图片数据输入到模型中
model.inputs["image"].set_value(data)
model.inputs["im_info"].set_value(im_info)
# 定义静态模型推理函数
from megengine import jit
@jit.trace(symbolic=True)
def infer():
predictions = model(model.inputs)
return predictions
# 推理预测
evaluator = models_api.DetEvaluator(model)
pred_res = evaluator.predict(infer)
# 可视化目标检测结果,生成标注后的图片
res_img = models_api.DetEvaluator.vis_det(
image, # 原始图片
pred_res, # 推理预测结果
is_show_label=True,
# classes=data_mapper[cfg.test_dataset["name"]].class_names,
)
cv2.imwrite("results1.jpg", res_img)
本仓库包含了采用MegEngine实现的各种主流深度学习模型。
official目录下提供了各种经典的图像分类、目标检测、图像分割以及自然语言模型的官方实现。每个模型同时提供了模型定义、推理以及训练的代码。
官方会一直维护official下的代码,保持适配MegEngine的最新API,提供最优的模型实现。同时,提供高质量的学习文档,帮助新手学习如何在MegEngine下训练自己的模型。
对于每个模型,提供了至少四个脚本文件:模型定义(model.py
)、模型推理(inference.py
)、模型训练(train.py
)、模型测试(test.py
)。
每个模型目录下都对应有一个README
,介绍了模型的详细信息,并详细描述了训练和测试的流程。例如 ResNet README。
另外,official
下定义的模型可以通过megengine.hub
来直接加载,如上案例代码。
更多可以通过megengine.hub
接口加载的模型见hubconf.py。
在开始运行本仓库下的代码之前,用户需要通过以下步骤来配置本地环境:
1.克隆仓库
git clone https://github.com/MegEngine/Models.git
2.安装依赖包
pip3 install --user -r requirements.txt
3.添加目录到python环境变量中
which python3
/home/megstudio/.miniconda/envs/xuan/bin/python3 # 我的python路径
/home/megstudio/workspace/Models # 我的Models路径
# 配置模型和python路径 export PYTHONPATH=/path/to/models:$PYTHONPATH
export PYTHONPATH="/home/megstudio/.miniconda/envs/xuan/bin/python3:/home/megstudio/workspace/Models"
# 中间用“:”隔开
# 查看配置成功
echo $PYTHONPATH
/home/megstudio/.miniconda/envs/xuan/bin/python3:/home/megstudio/workspace/Models
python3 workspace/Models/official/vision/detection/tools/inference.py \
-f workspace/Models/official/vision/detection/configs/retinanet_res50_coco_1x_800size.py \
-w ./workspace/retinanet_res50_coco_1x_800size.pkl \
-i workspace/cat/cat1.jpg
inference.py
的命令行选项如下:
-f
, 测试的网络结构描述文件。-i
, 需要测试的样例图片。-w
, 网络结构文件所对应的训练权重, 可以从顶部的表格中下载训练好的检测器权重。