Wie kann man im Tensorflow auf einen skalaren Tensorwert zugreifen, bevor er auf die GPU übertragen wird?

Question

In tensorflow registrieren I ein op etwa so:

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"); 

Alle oben genannten Eingänge sind Tensoren, aber ref_freq ein Skalar oder 0-D-Tensors. Im Compute() Methode meiner CPU kernel ich folgendes tun kann die skalare zu extrahieren:

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

die gleiche Art von Code erzeugt jedoch einen segfault im Compute() Methode meiner GPU Kernel , weil die CPU jetzt versucht, auf einen Speicherblock auf dem GPU-Gerät zuzugreifen. Gibt es einen Wert, der diesen skalaren Wert abfängt, bevor er in die GPU gesendet wird? Ich möchte vermeiden die folgende zusätzliche Ebene der Speicher indirection in einen CUDA Kernel:

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

Ich glaube nicht, Attr der Ansatz für ref_freq zu verwenden ist, da sie veränderbar ist, konfigurierbaren Wert.

1. CPU Tensorflow Kernelcode ist here.
2. GPU Tensorflow Kernel-Code ist here.
3. Python variable Setup-Code ist here

Answer 1

Sie können angeben, dass ein oder mehrere der Eingänge (oder Ausgänge) einer TensorFlow OpKernel sind in „Host-Speicher“, die man den Wert in der für den Zugriff ermöglicht Compute() Methode. Um dies zu tun, würden Sie Ihren REGISTER_KERNEL_BUILDER() Anruf ändern hinzufügen, um eine .HostMemory("ref_freq") Anweisung:

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>);