更改 Eigen::Tensor(Map) 的类型

Changing the type of a Eigen::Tensor(Map)

我目前正在构建自定义 TensorFlow Op。它应该像 Conv2D-Op 一样工作,除了它使用自定义数据类型。由于在 Eigen 中实现自定义数据类型相对容易,而在 TensorFlow 中实现起来相当困难,因此我决定在 TensorFlow 调用 Eigen 之前将 Eigen 张量复制到具有自定义数据类型的新 Eigen 张量。将 Eigen::TensorMap<Eigen::Tensor<float, 4, Eigen::RowMajor, Eigen::DenseIndex>, Eigen::Aligned> 转换为 Eigen::TensorMap<Eigen::Tensor<CustomType, 4, Eigen::RowMajor, Eigen::DenseIndex>, Eigen::Aligned>、运行 计算,然后再转换回 float

我在 TensorFlows conv_2d.h 中添加了一些代码,在 SpatialConvolutionoperator() 中。我写了两个辅助函数 convertToCustomTypeconvertFromCustomType 应该为我做转换。目前我不太关心性能。

所以基本上我在这一行之前和之后注入了我的两个转换函数:https://github.com/tensorflow/tensorflow/blob/r1.12/tensorflow/core/kernels/conv_2d.h#L72

template<typename T>
Eigen::Tensor<CustomType, 4, Eigen::RowMajor, Eigen::DenseIndex> convertToCustomType(T &input) {
    Eigen::Tensor<CustomType, 4, Eigen::RowMajor, Eigen::DenseIndex> ret;
    ret.resize(input.dimensions());

    for (int a = 0; a < ret.size(); a++) {
        ret(a) = input(a);
    }
    return ret;
}

template<typename T1, typename T2>
void convertFromCustomType(T1 &input, T2 &output) {
    for (int a = 0; a < input.size(); a++) {
        output(a) = input(a);
    }
}

template<typename Device, typename T>
struct SpatialConvolution {
    void operator()(const Device &d, typename TTypes<T, 4>::Tensor output,
                    typename TTypes<T, 4>::ConstTensor input,
                    typename TTypes<T, 4>::ConstTensor filter, int row_stride,
                    int col_stride, int row_dilation, int col_dilation,
                    const Eigen::PaddingType &padding) {


        auto input_c = convertToCustomType(input);
        auto filter_c = convertToCustomType(filter);
        auto output_c = convertToCustomType(output);

        SpatialConvolutionFunc(d, output_c, input_c, filter_c, row_stride, col_stride, row_dilation, col_dilation, padding);

        convertFromCustomType(output_approx, output);
        output.device(d) = output;
    }
};

我还尝试 运行 分别处理张量的所有 4 个维度,这似乎也不起作用。

template <typename T>
Eigen::Tensor<ApproxType, 4, Eigen::RowMajor> convertToCustomType(T input) {
    Eigen::Tensor<ApproxType, 4, Eigen::RowMajor> ret;
    ret.resize(input.dimensions());
    for (int a = 0; a < ret.dimension(0); a++) {
        for (int b = 0; b < ret.dimension(1); b++) {
            for (int c = 0; c < ret.dimension(2); c++) {
                for (int d = 0; d < ret.dimension(3); d++) {
                    ret(a, b, c, d) = input(a, b, c, d);
                }
            }
        }
    }
    return ret;
}

两个版本都可以编译,但产生的结果不正确。如果我 运行 我的整个 TensorFlow 网络与这个自定义 Op 它变得不确定,输出在不同的 运行 中变化,具有相同的输入。

0
[[-0.06590138]]
1
[[-0.04544453]]
2
[[-0.0286443]]
3
[[-0.06590138]]
4
[[-0.06590138]]
5
[[-0.04544453]]

我该如何更改特征张量的实际类型? 我注意到有一些像 Tensor::cast<T>() 这样优雅的东西,但是用 T 调用它是 Eigen::half 以外的任何东西都不会编译。我可能在我的自定义类型中遗漏了什么吗?

我知道这是一个非常具体的问题,但如果有任何想法,我将不胜感激。

显然创建张量是不够的,必须在填充之前用 ret.setZero() 初始化它。