c++ - TensorFlow自定义操作：输出张量的线程安全

Question

我有一个 TensorFlow 的自定义操作，它受益于并行执行特定计算。但是，一位用户在使用此并行计算时报告了崩溃（在我的计算机上一切正常）。

• 使用线程时，在 Compute 方法中有什么需要注意的吗？
• 读/写张量是否已经是线程安全的，还是我必须小心？
• 特别是，在写入输出张量时我是否必须同步我的线程？请参阅下面的简化代码片段。

    void Compute(OpKernelContext* context) 覆盖
    {
        // 输出张量
        张量* outputTensor = nullptr;
        OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape({maxI, maxJ}), &outputTensor));
        自动 outputMapped = outputTensor->tensor();

        // 对索引 (i, j) 进行一些计算，并将结果写入输出张量的索引 (i, j)
        // 对每个 (i, j) 组合执行此操作
        outputMapped(i, j) = 结果；

        // 并行计算每个索引 (i, j) 的结果将加快计算速度
        // 问题：是否允许不同步地从不同线程写入 outputMapped(i, j)？
    }

这是相关代码：

• Compute(...) 方法：输出张量
• 开始并行工作
• 并行写入每个线程的输出张量

我查看了TF 文档，但它只提到 TF 本身可能会并行调用 Compute(...) 。但我找不到任何有关张量线程安全的信息。有没有人有这方面的更多信息？

标签： c++multithreadingtensorflow