Entwerfen eines akkumulierenden Tensorflow-GPU-Operators

Question

Ich entwerfe einen GPU-Op-Kernel, der iterativ Daten in einem Puffer aus GPU-Speicher akkumuliert. Es ist wichtig, dass die Daten im GPU-Speicher verbleiben. So etwas entlang der Linien von:

with tf.device('/gpu:0'): 
    buffer = tf.zeros(...) 

    buffer = accumulate(param11, param12, buffer) 
    buffer = accumulate(param21, param22, buffer) 
    buffer = accumulate(param31, param32, buffer) 

with tf.device('/cpu:0'): 
    A = do_some_more_stuff(buffer) 

würde ich eine Eingabe auf drei Ansätze, dass ich denke, verwendet werden kann, um dies zu erreichen:

1. Ausgang Tensor bei jedem Aufruf zuweisen und verwenden, die als ein Eingangstensor beim nächsten Anruf. Dies ist einfach zu implementieren, aber ich bin besorgt, dass kontinuierliche Zuweisung von GPU-Speicher ein Problem sein wird. Wird Tensorflow jetzt ungenutzte Zuweisungen in den GPU-Speicherpool freigeben?

   REGISTER_OP("Accumulate") 
       .Input("param1: T") 
       .Input("param2: T") 
       .Input("buffer_in: T") 
       .Output("buffer_out: T") 
   
   void Compute(tensorflow::OpKernelContext * ctx) override 
   { 
       TensorShape output_shape{...}; 
       Tensor * output_ptr = nullptr; 
       OP_REQUIRES_OK(ctx, ctx->allocate_output(
        0, output_shape, &output_ptr)) 
   
       kernel<<<grid, blocks, 0, stream>>>(
        ctx->input(0), ctx->input(1), 
        output); 
   }  
   

2. Referenz Eingangs- und Ausgangs Tensoren und sicherzustellen, dass sie sich beziehen auf die gleichen Daten. Als ich die Standard-ops und OpKernelContext Dokumentation zu verstehen, muss dies mit einem Mutex als andere ops geschützt werden auch die zugrunde liegenden referenzierten Tensor zugreifen kann ...

   REGISTER_OP("Accumulate") 
       .Input("param1: T") 
       .Input("param2: T") 
       .Input("buffer_in: Ref(T)") 
       .Output("buffer_out: Ref(T)") 
   
   void Compute(tensorflow::OpKernelContext * ctx) override 
   { 
       mutex_lock(mu_); 
   
       ctx->forward_ref_input_to_ref_output(2, 0); 
   
       kernel<<<grid, blocks, 0, stream>>>(
        ctx->input(0), ctx->input(1), 
        ctx->mutable_input(2, true)); 
   } 
   

3. Verwenden allocate_persistent() in Verbindung mit einem OpKernelConstruction Kontext , um einen dauerhaften Puffer für die Akkumulation bereitzustellen. Ich würde es vorziehen, dies nicht zu tun, weil Ich habe mit variablen Puffergrößen zu tun, und sie werden wahrscheinlich ziemlich groß sein.

Answer 1

Ich bin nicht wirklich sicher, was Sie versuchen, mit Ihrem C++ Code zu tun, aber in dem Python-Schnipsel aus der Suche denke ich tf.assign helfen könnte. Es erlaubt Ihnen, solche Dinge zu tun:

buffer = tf.Variable(...) 
param = tf.Variable(...) 
accumulate_op = buffer.assign(expr<param, buffer>) 

... 

sess.run(accumulate_op) 

Lauf accumulate_op sollte Ihr Puffer auf dem gpu aktualisieren (Sie können es in einem tf.group müssen wickeln den aktualisierten Wert zu vermeiden holen).