我正面临对象检测问题。该对象包含2个组件，并带有2个矩形边界框。我熟悉基于深度学习的对象检测算法，例如RetinaNet等。
因此，我的想法是对每个组件检测使用RetinaNet。使用RetinaNet可以轻松实现这一点，在此阶段之后，该模型将输出4个张量：o1-box-回归，o1-label-分类，o2-box-回归和o2-label-回归。在这一层的顶部，我需要组合这两个组件。
由于对象必须包含一个o1和一个o2，因此我需要为每个o1找到一个最佳o2。换句话说，我需要选择在某些约束下具有max（o2-label-classification）的o2（例如o2离o1不会很远）。
在我看来，我将不得不使用arg_max操作来找到索引，然后使用arg_max返回的索引张量来选择o2-label-regression和o2-label-classification。但是，我注意到arg_max运算不可微，因此无法计算用于反向传播的梯度，这基本上使该模型无法使用梯度下降方法进行训练。
任何人都可以就如何解决此问题分享一些想法吗？确实非常感谢。
顺便说一句，我可以分别为o1和o2训练2个RetinaNet。然后建立另一个合并模型（无论是否可训练），以在第二步中将它们合并。但是，此刻，我正在探索建立具有2个组件的端到端可训练模型进行对象检测的可能性。