基于华为atlas的皮带跑偏、空载、堆煤、启停探索

生乎吾前，其闻道也固先乎吾，吾从而师之；生乎吾后，其闻道也亦先乎吾，吾从而师之。吾师道也，夫庸知其年之先后生于吾乎？是故无贵无贱，无长无少，道之所存，师之所存也。

写在前面：

写这篇的时候，想起当年第一次接触atlas还是在京东的一次aicon的会议上，其实那时觉得这东西挺新的，还有自己的IDE，其实自己也没用过。后来的工作中在面试中，也见过1-2个用过atlas的面试者。直到几年前归梓，阴差阳错，华为atlas满足信创的需求，才于23年12月正式研究学习atlas。几个月过去了，整体感觉昇腾生态闭塞、最新的git很难支持、华为各种软件包下载也很坑。问题是很多的，至于其他国产显卡，像曙光、海光、摩尔线程、寒武纪等自己也没机会接触，不得不说华为在信创这条路的部署是强于其他厂商的，于是我便这条路走下去了。

回归主题：

在真正做这件事情的时候，其实也预研了几天，也了解了别的公司的大体思路，比如精英、壹合等。其实我这里想做一个不一样的，有别于别人的，视觉有冲击力的，真正的all in one模型。

整体模型我没有选择基于检测的思路，而是选择基于分割的思路，因为分割可以帮我获取更多的检测不具备的信息，这里后续可以看到。皮带跑偏、空载、堆煤基于unet实现，皮带启停基于unet+光流跟踪实现。

整体感觉模型这块不算复杂，唯一的麻烦的地方就是皮带、煤、煤块这几个的分割是属于多标签分割问题，就是说一个像素可以属于其中的一个也可以是属于其中的几个。这块需要在损失函数处下一些功夫修改代码，当然也有其他曲径通幽的方法。这个修改不能算复杂，复杂的是这个修改同时还得兼容atlas的模型转化，这块就需要真正下功夫了。

另外一个复杂的地方就是皮带中轴线的方法，探索这个中轴线大概耗费了1个月的时间把，中间尝试了很多方法，比如label_centerlines、PCA、旋转角度法、曲线拟合等，效果都不太满意。这里需要考虑皮带多种形状的自适应性，还得考虑效果、考虑速度。最终根据自己的思想实现了一个中轴线确定的方法，客观的说该方法不是完美的，但是已经可以解决90%以上的场景，具体后面我会有图片展示。

数据制作：

这里我一共标注了1056张图片，数据的标签包括，左托锟、右托锟、皮带、煤、煤堆、背景一共6个类别，标注工具选择labelme。

模型训练：

原型代码链接：https://github.com/milesial/Pytorch-UNet.git

训练过程采用累进训练的方式，我是一个一个目标递进训练的，这样可以获取更好的精度，具体的先训练出背景、皮带、左右托锟的模型，再在此基础上迭代煤的模型，最后迭代煤块的模型。整体感受随着类别数量的增加，训练过程会逐渐变得困难，尤其到最后一个目标得时候，如果采用一起训练得方法，精度会比我这种累进训练得方法至少低5个点。

最开始得时候我训练得模型，输入分辨率是572*572的，经过测试速度很慢，后来修改为384*384的，同时注意原始代码中的上采用部分要做适当的修改。这样才能保证上下采样都是2的整数倍。这里我直接简单粗暴的去掉了Up模块中的pad部分。

Atlas模型转化：

atc --model=./unet.onnx --framework=5 --output=unet --soc_version=Ascend310P3  --input_shape="input.1:1,3,384,384" --output_type="/outc/conv/Conv:0:FP32" --out_nodes="/outc/conv/Conv:0"

其中fusion_result.json内容，

[{"graph_fusion": {"ConstToAttrPass": {"effect_times": "4","match_times": "4"},"Conv2DbpInputDilationFusionPass": {"effect_times": "0","match_times": "4"},"ConvConcatFusionPass": {"effect_times": "0","match_times": "4"},"ConvFormatRefreshFusionPass": {"effect_times": "0","match_times": "19"},"ConvToFullyConnectionFusionPass": {"effect_times": "0","match_times": "19"},"ConvWeightCompressFusionPass": {"effect_times": "0","match_times": "19"},"CubeTransFixpipeFusionPass": {"effect_times": "0","match_times": "1"},"DeconvWeightTransFusionPass": {"effect_times": "0","match_times": "4"},"FIXPIPEAPREQUANTFUSIONPASS": {"effect_times": "0","match_times": "23"},"FIXPIPEFUSIONPASS": {"effect_times": "0","match_times": "23"},"RefreshInt64ToInt32FusionPass": {"effect_times": "1","match_times": "1"},"RemoveCastFusionPass": {"effect_times": "0","match_times": "48"},"SameInputConv2dPass": {"effect_times": "0","match_times": "10"},"SplitConvConcatFusionPass": {"effect_times": "0","match_times": "4"},"StrideHoistingPass": {"effect_times": "0","match_times": "18"},"TransdataCastFusionPass": {"effect_times": "0","match_times": "25"},"ZConcatDFusionPass": {"effect_times": "0","match_times": "4"}},"session_and_graph_id": "0_0","ub_fusion": {"TbeConvCommonRules0FusionPass": {"effect_times": "18","match_times": "18","repository_hit_times": "0"}}
}]

ge_check_op.json内容，

{"graph_id": 0,"op": [{"error_type": "infer_shape_error","input0": {"data_type": "DT_FLOAT","layout": "ND","shape": [1,128,80,80]},"input1": {"data_type": "DT_FLOAT","layout": "ND","shape": [1,128,81,81]},"name": "/up1/Concat_1","output0": {"data_type": "DT_FLOAT","layout": "ND","shape": [1,128,80,80]},"reason": "InferShapeFailed!","type": "ConcatD"}],"session_id": 0
}