テンソルフローでは、GPUに移動する前にスカラーテンソル値にどのようにアクセスしますか？ tensorflowで

Question

IはそうのようなOP登録：

REGISTER_OP("RimeBSqrt") 
    .Input("stokes: FT") 
    .Input("alpha: FT") 
    .Input("frequency: FT") 
    .Input("ref_freq: FT") 
    .Output("b_sqrt: CT") 
    .Attr("FT: {float, double} = DT_FLOAT") 
    .Attr("CT: {complex64, complex128} = DT_COMPLEX64"); 

上記入力のすべてが、テンソルである しかしref_freqはスカラーまたは0次元テンソルです。

const Tensor & in_ref_freq = context->input(3); 
FT ref_freq = in_ref_freq.tensor<FT, 1>()(0); 

ただし、同じ種類のコードは、私のGPUカーネルの計算（）メソッドでセグメンテーションフォルト を生成します。私はスカラーを抽出するために、次の操作を行うことができ、私のCPUのカーネル の計算（）メソッドで CPUが GPUデバイスのメモリブロックにアクセスしようとしているためです。とにかくこのスカラーをGPUに送る前に値 を傍受することはありますか？私はAttrはそれを変更、設定可能な値であるため、ref_freqに使用するアプローチであるとは思わない

template <typename FT> 
__global__ void kernel(..., FT * ref_freq, ...) 
{ 
    FT value = ref_freq[0]; 
} 

：私は でCUDAカーネルをメモリ間接の 以下の余分なレベルを避けるためにしたいと思います。

1. CPU Tensorflowカーネルコードはhereです。
2. GPU Tensorflowカーネルコードはhereです。
3. Pythonの変数の設定コードを使用すると、1つまたは複数の入力のへ（または出力から）TensorFlow OpKernelはあなたが値にアクセスすることを可能にする「ホストメモリ」、であることを指定することができhere

Answer 1

ですCompute()方法。

REGISTER_KERNEL_BUILDER(
    Name("RimeBSqrt") 
    .Device(tensorflow::DEVICE_GPU) 
    .TypeConstraint<float>("FT") 
    .TypeConstraint<tensorflow::complex64>("CT") 
    .HostMemory("ref_freq"), 
    RimeBSqrt<tensorflow::GPUDevice, float, tensorflow::complex64>);