详细记录YOLACT实例分割ncnn实现

作者:nihui

链接:https://zhuanlan.zhihu.com/p/128974102

本文转载自知乎,作者已授权,未经许可请勿二次转载。

0x0 YOLACT实例分割

https://urlify.cn/rURFry

  • 端到端一阶段完成实例分割

  • 速度快,550x550图片在TitanXP上号称达到33FPS

  • 开源代码,pytorch大法好!

0x1 缘由

纵观整个github,无论是ncnn还是ncnn衍生项目,分类,检测,定位,特征提取,OCR,风格转换....

然而,就是没有找到实例分割的例子,以至于有人发了个issue,并点名要求搞个 YOLACT 实例分割 https://github.com/Tencent/ncnn/issues/1679

那好吧于是写个YOLACT例子,顺带介绍下如何用ncnn实现类似这种需要后处理的算法

0x2 pytorch测试

YOLACT项目里有YOLACT++模型,速度更快,效果更好,不过YOLACT++用了个对部署不友好的经典骚操作deformable convolution

假装没看到,我们去下载YOLACT模型

新建weights文件夹,下载 yolact_resnet50_54_800000.pth

根据 README 指示,先拿张图试试看效果

$ python eval.py --trained_model=weights/yolact_resnet50_54_800000.pth --score_threshold=0.15 --top_k=15 --image=test.jpg

0x3 去掉后处理导出onnx

直接修改 eval.py 的 evalimage,把结果展示换成 onnx export

def evalimage(net:Yolact, path:str, save_path:str=None):
    frame = torch.from_numpy(cv2.imread(path)).cuda().float()
    batch = FastBaseTransform()(frame.unsqueeze(0))
    preds = net(batch)

    torch.onnx._export(net, batch, "yolact.onnx", export_params=True, keep_initializers_as_inputs=True, opset_version=11)

根据YOLACT issue中的信息,yolact.py开头的JIT要关掉才能导出onnx

# As of March 10, 2019, Pytorch DataParallel still doesn't support JIT Script Modules
use_jit = False

YOLACT后处理部分写得非常 pythonic,这样直接导出不行,要把后处理从模型剔除,方便导出转换

即便onnx能导出后处理,也不建议这么做

  • 后处理部分没有标准化,每个项目作者的实现细节也各不相同,比如各种nms和bbox计算方式,ncnn很难用统一的op实现(caffe-ssd因为只有一种版本,所以有实现)

  • 后处理在onnx中会转换成一大坨胶水op,非常琐碎,在框架中实现效率低下

  • onnx的大部分胶水op,ncnn不支持或有兼容问题,比如Gather等,无法直接使用

因此,去掉后处理导出onnx,是正确转换 pytorch ssd 等类似模型的通常做法

打开yolact.py,找到 class Yolact 的 forward 方法,把 detect 过程去掉,直接返回模型的 pred_outs 输出

            # return self.detect(pred_outs, self)
            return pred_outs;

再一次跑一遍图片测试,不包含后处理的 yolact.onnx 出现了

$ python eval.py --trained_model=weights/yolact_resnet50_54_800000.pth --score_threshold=0.15 --top_k=15 --image=test.jpg

0x4 简化onnx

直接导出的onnx模型有很多胶水op是ncnn不支持的,用onnx-simplifier是常规操作

$ pip install -U onnx --user
$ pip install -U onnxruntime --user
$ pip install -U onnx-simplifier --user

$ python -m onnxsim yolact.onnx yolact-sim.onnx

这时候遇到个问题

Graph must be in single static assignment (SSA) form, however '523' has been used as output names multiple times

经过在github翻看issue,确认这是 onnx bug

https://link.zhihu.com/?target=https%3A//github.com/onnx/onnx/issues/2613

幸好 onnx-simplifier 已提供办法绕过

$ python -m onnxsim --skip-fuse-bn yolact.onnx yolact-sim.onnx

0x5 ncnn模型转换和优化

前面简化onnx的时候,--skip-fuse-bn 跳过了 batchnorm 合并,不过没关系,ncnn 也有这个功能

ncnnoptimize 工具实现了很多种算子融合,比如常见的 convolution-batchnorm-relu 等等

最后的参数 0 表示fp32模型,65536 表示精简为fp16模型,能减少模型二进制体积

$ ./onnx2ncnn yolact-sim.onnx yolact.param yolact.bin
$ ./ncnnoptimize yolact.param yolact.bin yolact-opt.param yolact-opt.bin 0

0x6 手工微调模型

还是这句话,不报错不代表一定能用,先用netron工具打开param看看模型结构

这个模型输出有四个,用红框框出来了

Convolution              Conv_263                 1 1 617 619 0=32 1=1 5=1 6=8192 9=1
Permute                  Transpose_265            1 1 619 620 0=3
UnaryOp                  Tanh_400                 1 1 814 815 0=16
Concat                   Concat_401               5 1 634 673 712 751 790 816 0=-3
Concat                   Concat_402               5 1 646 685 724 763 802 817 0=-3
Concat                   Concat_403               5 1 659 698 737 776 815 818 0=-3
Softmax                  Softmax_405              1 1 817 820 0=1 1=1

YOLACT 的后处理需要 loc conf prior mask maskdim 这些东西

一开始看不出这几个输出对应的是什么,那么就先看shape

ncnn::Extractor ex = yolact.create_extractor();

ncnn::Mat in(550, 550, 3);
ex.input("input.1", in);

ncnn::Mat b620;
ncnn::Mat b816;
ncnn::Mat b818;
ncnn::Mat b820;
ex.extract("620", b620);// 32 x 138x138
ex.extract("816", b816);// 4 x 19248
ex.extract("818", b818);// 32 x 19248
ex.extract("820", b820);// 81 x 19248

直接编译运行发现 Concat 层 crash,即图中蓝框,Concat axis 参数是负数 0=-3,ncnn 还不支持

根据 Concat 多个输入shape,发现是二维数据在 h axis concat,直接改成 0=0 就可以替代

Concat                   Concat_401               5 1 634 673 712 751 790 816 0=0
Concat                   Concat_402               5 1 646 685 724 763 802 817 0=0
Concat                   Concat_403               5 1 659 698 737 776 815 818 0=0

b820在softmax后面,确信是 conf,shape 81x19248 表示 81分类 x 19248个prior

b816 shape 4x19248,对应于每个priorbox的bbox的偏移值

b818 shape 32x19248,根据YOLACT的后处理看,表示的是 maskdim,即32个分割热图的系数

b620 shape 32x138x138,即32个分割热图,前面有个permute层是NCHW->NHWC的转换 prior没有在模型中输出

ncnn 处理 b620 NHWC shape 不方便,改为 extract permute 前的 NCHW 数据 b619,即图中绿框输出

ncnn::Extractor ex = yolact.create_extractor();

ncnn::Mat in(550, 550, 3);
ex.input("input.1", in);

ncnn::Mat maskmaps;
ncnn::Mat location;
ncnn::Mat mask;
ncnn::Mat confidence;
ex.extract("619", maskmaps);// 138x138 x 32
ex.extract("816", location);// 4 x 19248
ex.extract("818", mask);// maskdim 32 x 19248
ex.extract("820", confidence);// 81 x 19248

0x7 生成prior

原始代码在 yolact.py class PredictionModule make_priors,增加一些 print 获得全部 priorbox 生成规则超参

const int conv_ws[5] = {69, 35, 18, 9, 5};
const int conv_hs[5] = {69, 35, 18, 9, 5};

const float aspect_ratios[3] = {1.f, 0.5f, 2.f};
const float scales[5] = {24.f, 48.f, 96.f, 192.f, 384.f};

YOLACT的prior四个数值是 center_x center_y box_w box_h,值域 0~1

作者当初写了个bug,box_h = box_w 固定成方的了,我们也要把这个bug复现出来

// make priorbox
ncnn::Mat priorbox(4, 19248);
{
    float* pb = priorbox;

    for (int p = 0; p < 5; p++)
    {
        int conv_w = conv_ws[p];
        int conv_h = conv_hs[p];

        float scale = scales[p];

        for (int i = 0; i < conv_h; i++)
        {
            for (int j = 0; j < conv_w; j++)
            {
                // +0.5 because priors are in center-size notation
                float cx = (j + 0.5f) / conv_w;
                float cy = (i + 0.5f) / conv_h;

                for (int k = 0; k < 3; k++)
                {
                    float ar = aspect_ratios[k];

                    ar = sqrt(ar);

                    float w = scale * ar / 550;
                    float h = scale / ar / 550;

                    // This is for backward compatability with a bug where I made everything square by accident
                    // cfg.backbone.use_square_anchors:
                    h = w;

                    pb[0] = cx;
                    pb[1] = cy;
                    pb[2] = w;
                    pb[3] = h;

                    pb += 4;
                }
            }
        }
    }
}

0x8 YOLACT全流程实现

预处理部分

data/config.py 有 ImageNet 的 MEAN STD,BGR顺序

# These are in BGR and are for ImageNet
MEANS = (103.94, 116.78, 123.68)
STD   = (57.38, 57.12, 58.40)

YOLACT实际输入RGB,要换下顺序

const int target_size = 550;

int img_w = bgr.cols;
int img_h = bgr.rows;

ncnn::Mat in = ncnn::Mat::from_pixels_resize(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, img_w, img_h, target_size, target_size);

const float mean_vals[3] = {123.68f, 116.78f, 103.94f};
const float norm_vals[3] = {1.0/58.40f, 1.0/57.12f, 1.0/57.38f};
in.substract_mean_normalize(mean_vals, norm_vals);

后处理部分

这部分和 SSD 后处理非常类似,sort nms 这些代码抠 ncnn/src/layer/detectionoutput.cpp

唯一要注意的地方就是 bbox 生成和 SSD 不一样,要用 center_x center_y box_w box_h 实现,YOLACT原代码在 layers/box_util.py decode 函数

YOLACT有fastnms方法 layers/funstions/detection.py,速度更快,可我觉得普通nms毕竟是现成代码,用着挺好的

// generate all candidates for each class
for (int i=0; i<num_priors; i++)
{
    // find class id with highest score
    // start from 1 to skip background

    // ignore background or low score
    if (label == 0 || score <= confidence_thresh)
        continue;

    // apply center_size to priorbox with loc
    float var[4] = {0.1f, 0.1f, 0.2f, 0.2f};

    float pb_cx = pb[0];
    float pb_cy = pb[1];
    float pb_w = pb[2];
    float pb_h = pb[3];

    float bbox_cx = var[0] * loc[0] * pb_w + pb_cx;
    float bbox_cy = var[1] * loc[1] * pb_h + pb_cy;
    float bbox_w = (float)(exp(var[2] * loc[2]) * pb_w);
    float bbox_h = (float)(exp(var[3] * loc[3]) * pb_h);

    float obj_x1 = bbox_cx - bbox_w * 0.5f;
    float obj_y1 = bbox_cy - bbox_h * 0.5f;
    float obj_x2 = bbox_cx + bbox_w * 0.5f;
    float obj_y2 = bbox_cy + bbox_h * 0.5f;

    // clip inside image

    // append object candidate
}

// merge candidate box for each class
for (int i=0; i<(int)class_candidates.size(); i++)
{
    // sort + nms
}

// sort all result by score

// keep_top_k

分割图生成

maskmaps 实际是 32 张 138x138 尺寸的热图,前面输出的每个 object 都自带 32 个 float 系数

object 的分割图就是每张热图 * 对应系数,求和,放大到原图尺寸,二值化,最后 crop inside 输出框

unnatrual很好看的!

0x9 补充学习资料

噫?还有补充学习资料?

ncnn实现代码和转好的模型已上传到github

https://link.zhihu.com/?target=https%3A//github.com/Tencent/ncnn

☆ END ☆

如果看到这里,说明你喜欢这篇文章,请转发、点赞。微信搜索「uncle_pn」,欢迎添加小编微信「 mthler」,每日朋友圈更新一篇高质量博文(无广告)。

扫描二维码添加小编↓

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值