Заставьте TensorFlow использовать данные обучения, сгенерированные на лету с помощью пользовательской процедуры CUDA

Предположим, что мы генерируем наши собственные обучающие данные (путем выборки из некоторого процесса распространения и вычисления, например, некоторого количества интересующего его количества) и что у нас есть собственная подпрограмма CUDA под названием generate_data, которая генерирует метки в памяти графического процессора для данного набора входных данных.

Следовательно, мы находимся в специальной настройке, где мы можем генерировать столько пакетов обучающих данных, сколько мы хотим, в режиме "онлайн" (на каждой итерации пакета мы вызываем эту процедуру generate_data, чтобы сгенерировать новый пакет и отбросить старый пакет).

Поскольку данные генерируются на графическом процессоре, есть ли способ заставить TensorFlow (API-интерфейс Python) напрямую использовать его в процессе обучения? (например, для заполнения заполнителя). Таким образом, такой конвейер будет эффективным.

Насколько я понимаю, в настоящее время вам потребуется такая настройка, чтобы скопировать данные из графического процессора в процессор, а затем позволить TensorFlow снова скопировать их из процессора в графический процессор, что довольно расточительно, поскольку выполняются ненужные копии.

РЕДАКТИРОВАТЬ: если это помогает, мы можем предположить, что подпрограмма CUDA реализована с использованием JIT-компилятора Numba CUDA.

Ответы

Ответ 1

Это определенно не полный ответ, но, надеюсь, может помочь.

  • Вы можете интегрировать свою подпрограмму CUDA в TensorFlow, написав собственный оператор. В настоящее время в TensorFlow нет другого способа взаимодействия с другими подпрограммами CUDA.

  • Что касается написания учебного цикла полностью на GPU, мы можем написать подпрограмму на GPU, используя tf.while_loop, очень похоже на этот вопрос SO:

    i = tf.Variable(0, name='loop_i')
    
    def cond(i):
        return i < n
    
    def body(i):
        # Building the graph for custom routine and our model
        x, ground_truth = CustomCUDARountine(random_seed, ...)
        predictions = MyModel(x, ...)
    
        # Defining the optimizer
        loss = loss_func(ground_truth, predictions)
        optim = tf.train.GradientDescentOptimizer().minimize(loss)
    
        # loop body
        return tf.tuple([tf.add(i, 1)], control_inputs=[optim])
    
    loop = tf.while_loop(cond, body, [i])
    
    # Run the loop
    tf.get_default_session().run(loop)