如何在Tensorflow中覆盖默认的conv内核并使用自定义库?

时间:2019-06-12 06:48:21

标签: tensorflow

我正在与其他一些操作一起针对项目优化我的机器的自定义卷积例程,并寻求将其与Tensorflow链接。这与MKLDNN的意图相似,但是规模要小得多。我看到MKLDNN很好地解耦了,因为它是一个独立的库,从源代码构建时可以链接到Tensorflow。我也想在我的图书馆中重现这一点。

我已经检查了官方的TF页面,在这里找到了这个页面-https://www.tensorflow.org/guide/extend/op。这对于添加新操作很有帮助,但是这不是我的意图。我寻求替换现有的op,例如conv2D。更具体地说,我想将调用从接口重定向到我的实现。

hack可能是直接在TF中直接修改相应的内核源代码。但是,这显然是微不足道的,并且不应该共享。我猜想会有更好的解决方案,MKLDNN的成功使这种信念更加坚定。

我们非常感谢您的协助!谢谢。

0 个答案:

没有答案