首页 > 解决方案 > 使用自定义 oneDNN 库实现构建 TensorFlow 的步骤是什么?

问题描述

我正在使用一个自定义 oneDNN 库实现,我需要 Tensorflow (v2.4.0) 来构建它。但是,我注意到在构建 Tensorflow 时没有使用系统提供的 OneDNN 库的构建选项。我想知道支持使用系统提供的 oneDNN 库的 Tensorflow 构建的步骤是什么。

一些具体...

  1. oneDNN 库版本为 1.6.4,已安装在系统中(Linux)。此版本对应于使用“--config=mkl_opensource_only”Bazel 标志编译时使用的 Tensforflow。
  2. 我可以访问库源代码,但最好使用已编译的库。
  3. 目标架构是 RISC-V,操作系统是 Linux。

标签: tensorflowbuildbazelintel-mkl

解决方案


没有简单的方法告诉 bazel 链接到自定义库,但如果您有修改后的 oneDNN 源目录,则可以编辑tensorflow/workspace.bzl文件并用mkl_dnn_v1new_local_repository 规则替换存储库定义以指向修改后的源目录。即替换这个块:

    tf_http_archive(
        name = "mkl_dnn_v1",
        build_file = clean_dep("//third_party/mkl_dnn:mkldnn_v1.BUILD"),
        sha256 = "5369f7b2f0b52b40890da50c0632c3a5d1082d98325d0f2bff125d19d0dcaa1d",
        strip_prefix = "oneDNN-1.6.4",
        urls = [
            "https://storage.googleapis.com/mirror.tensorflow.org/github.com/oneapi-src/oneDNN/archive/v1.6.4.tar.gz",
            "https://github.com/oneapi-src/oneDNN/archive/v1.6.4.tar.gz",
        ],
    )

像这样:

    native.new_local_repository(
        name = "mkl_dnn_v1",
        build_file = clean_dep("//third_party/mkl_dnn:mkldnn_v1.BUILD"),
        path = "/path/to/your/modified/oneDNN/sources",
    )

如果您添加了任何新的源文件,您可能还想修改third_party/mkl_dnn/mkldnn_v1.BUILD文件。

附录:

--config=mkl_opensource_only现在似乎坏了,您可能会使用 just--config=mkl来获得更好的运气。


推荐阅读