Почему выбрана реализация CPU моего пользовательского op?

чтобы узнать, как писать пользовательские TensorFlow ops, я следовал добавление нового Op учебник и сделал" add_b " op, который добавляет скаляр b для каждого входного значения.

add_b_op.cc:

#define EIGEN_USE_THREADS

#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"

#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"

using namespace tensorflow;

REGISTER_OP("AddB")
    .Attr("T: {float, double}")
    .Input("input: T")
    .Input("b: T")
    .Output("output: T")
    .SetShapeFn([] (shape_inference::InferenceContext* c) -> Status {
      shape_inference::ShapeHandle out;
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &out));
      return shape_inference::UnchangedShape(c);
    })
//----------------------------------------------------------------------
    .Doc(R"doc(
Adds `b` to each input.

input: The input values.
b: A number to add to each input value.
)doc");


template <typename T>
class AddBCpuOp : public OpKernel {
 public:
  explicit AddBCpuOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    const Tensor& input_tensor = context->input(0);
    const auto input = input_tensor.flat<T>();

    Tensor* output_tensor = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));
    auto output = output_tensor->flat<T>();

    const Eigen::ThreadPoolDevice& d = context->eigen_device<Eigen::ThreadPoolDevice>();

    // Note: The mistake of adding 1 instead of `b` is intentional to be able to distinguish
    // the CPU and GPU implementations.
    output.device(d) = input + static_cast<T>(1);
  }
};

REGISTER_KERNEL_BUILDER(
    Name("AddB")
    .Device(DEVICE_CPU)
    .TypeConstraint<float>("T"),
    AddBCpuOp<float>);
REGISTER_KERNEL_BUILDER(
    Name("AddB")
    .Device(DEVICE_CPU)
    .TypeConstraint<double>("T"),
    AddBCpuOp<double>);


#if GOOGLE_CUDA

template <typename T>
bool LaunchAddBKernel(const T *__restrict__ d_input, int n, const T *__restrict__ d_b, T *__restrict__ d_output);

template <typename T>
class AddBGpuOp : public OpKernel {
 public:
  explicit AddBGpuOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    const Tensor& input_tensor = context->input(0);
    const auto input = input_tensor.flat<T>();

    const Tensor& b_tensor = context->input(1);
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(b_tensor.shape()),
                errors::InvalidArgument("add_b expects a scalar for `b`."));
    const auto b = b_tensor.scalar<T>();

    Tensor* output_tensor = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));
    auto output = output_tensor->flat<T>();

    OP_REQUIRES(context, LaunchAddBKernel(input.data(), input.dimension(0), b.data(), output.data()),
                errors::Internal("add_b: LaunchAddBKernel() failed."));
  }
};

REGISTER_KERNEL_BUILDER(
    Name("AddB")
    .Device(DEVICE_GPU)
    .TypeConstraint<float>("T"),
    AddBGpuOp<float>);
REGISTER_KERNEL_BUILDER(
    Name("AddB")
    .Device(DEVICE_GPU)
    .TypeConstraint<double>("T"),
    AddBGpuOp<double>);

#endif // if GOOGLE_CUDA

add_b_op.cu.cc

template <typename T, int BLOCK_DIM_X>
__global__ void AddBKernel(const T *__restrict__ d_input, int n, const T *__restrict__ d_b, T *__restrict__ d_output) {
  const int i = blockIdx.x * BLOCK_DIM_X + threadIdx.x;
  if (i < n) {
    d_output[i] = d_input[i] + *d_b;
  }
}

template <typename T>
bool LaunchAddBKernel(const T *__restrict__ d_input, int n, const T *__restrict__ d_b, T *__restrict__ d_output) {
  if (n <= 0) return true;

  constexpr int BLOCK_DIM_X = 256;
  AddBKernel<T, BLOCK_DIM_X><<<n / BLOCK_DIM_X + (n % BLOCK_DIM_X != 0), BLOCK_DIM_X>>>(d_input, n, d_b, d_output);
  return true;
}

// Explicit instantiations.
template bool LaunchAddBKernel<float>(const float *__restrict__, int, const float *__restrict__, float *__restrict__);
template bool LaunchAddBKernel<double>(const double *__restrict__, int, const double *__restrict__, double *__restrict__);

я намеренно ввел ошибку в реализации CPU, чтобы иметь возможность различать, используется ли реализация CPU или GPU.

когда я тестирую свой пользовательский op с:

from __future__ import print_function
import tensorflow as tf

module = tf.load_op_library('custom_ops.so')
with tf.Session(config = tf.ConfigProto(log_device_placement = True)):
  print(module.add_b([5., 4., 3., 2., 1.], 8.).eval())

я получаю следующий вывод:

I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:892] OS X does not support NUMA - returning NUMA node zero
I tensorflow/core/common_runtime/gpu/gpu_device.cc:951] Found device 0 with properties: 
name: GeForce GT 750M
major: 3 minor: 0 memoryClockRate (GHz) 0.9255
pciBusID 0000:01:00.0
Total memory: 2.00GiB
Free memory: 1.80GiB
I tensorflow/core/common_runtime/gpu/gpu_device.cc:972] DMA: 0 
I tensorflow/core/common_runtime/gpu/gpu_device.cc:982] 0:   Y 
I tensorflow/core/common_runtime/gpu/gpu_device.cc:1041] Creating TensorFlow device (/gpu:0) -> (device: 0, name: GeForce GT 750M, pci bus id: 0000:01:00.0)
Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: GeForce GT 750M, pci bus id: 0000:01:00.0
I tensorflow/core/common_runtime/direct_session.cc:252] Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: GeForce GT 750M, pci bus id: 0000:01:00.0

AddB: /job:localhost/replica:0/task:0/gpu:0
I tensorflow/core/common_runtime/simple_placer.cc:819] AddB: /job:localhost/replica:0/task:0/gpu:0
AddB/b: /job:localhost/replica:0/task:0/gpu:0
I tensorflow/core/common_runtime/simple_placer.cc:819] AddB/b: /job:localhost/replica:0/task:0/gpu:0
AddB/input: /job:localhost/replica:0/task:0/gpu:0
I tensorflow/core/common_runtime/simple_placer.cc:819] AddB/input: /job:localhost/replica:0/task:0/gpu:0
[ 6.  5.  4.  3.  2.]

"журналы размещения устройств", по-видимому, указывают, что op выполняется на GPU, но вывод указывает, что используется реализация CPU.

когда я комментирую две регистрации REGISTER_KERNEL_BUILDER () для DEVICE_CPU реализация, перекомпиляция и повторное тестирование, я получаю ожидаемый результат [ 13. 12. 11. 10. 9.], но есть ошибка:

E tensorflow/core/common_runtime/executor.cc:334] Executor failed to create kernel. Not found: No registered 'AddB' OpKernel for CPU devices compatible with node AddB = AddB[T=DT_FLOAT, _device="/job:localhost/replica:0/task:0/gpu:0"](AddB/input, AddB/b)
    .  Registered:  device='GPU'; T in [DT_FLOAT]
  device='GPU'; T in [DT_DOUBLE]

     [[Node: AddB = AddB[T=DT_FLOAT, _device="/job:localhost/replica:0/task:0/gpu:0"](AddB/input, AddB/b)]]

эту ошибку сообщение выглядит как ошибка для меня, потому что, хотя ошибка говорит "Executor не удалось создать ядро", ядро, по-видимому, было создано для запуска op на GPU.

почему используется реализация CPU, а не реализация GPU?

в случае, если это важно, вот подробности о моей настройке разработки:

  • я использую MacBook Pro со встроенным NVIDIA GeForce GT 750M (CUDA Compute Capability 3.0).
  • в macOS Сьерра-версия 10.12.1 (16B2555)
  • cuda_8.0.47_mac, cudnn-8.0-osx-x64-v5.1
  • TensorFlow 0.11.0rc2 установлен через:export TF_BINARY_URL=https://storage.googleapis.com/tensorflow/mac/gpu/tensorflow-0.11.0rc2-py2-none-any.whl

обновление я обнаружил, что выбор реализации CPU или GPU зависит от размера входного сигнала. Используя этот тестовый скрипт:

from __future__ import print_function
import numpy as np
import tensorflow as tf
from time import time

NUM_VALUES = 1310720

input = np.arange(0, NUM_VALUES, dtype = float)

module = tf.load_op_library('custom_ops.so')
with tf.Session(config = tf.ConfigProto(log_device_placement = True)):
  start = time(); print(module.add_b(input, 8.).eval()); end = time(); print(end - start)

.. когда NUM_VALUES 1310720 или меньше, тогда используется реализация CPU. Когда NUM_VALUES является 1310721 или более, затем GPU используется реализация.

есть ли (1310720 * 8 байт на двойной = ) 10 MiB отсечка? Если да, то как я могу переопределить это? Операция AddB () достаточно проста, но для более сложной пользовательской операции 10 MiB может быть слишком большим порогом для выбора реализации GPU.

3 ответов


Я думаю, что экземпляры шаблонов могут быть неправильными:

template <typename Device, typename T>
class AddBOp : public OpKernel {
...
}

REGISTER_KERNEL_BUILDER(
    Name("AddB")
    .Device(DEVICE_CPU)
    .TypeConstraint<float>("T"),
    AddBOp<CPUDevice, float>);

и затем:

template <typename T>
class AddBOp<GPUDevice, T> : public OpKernel {
...
}

REGISTER_KERNEL_BUILDER(
    Name("AddB")
    .Device(DEVICE_GPU)
    .TypeConstraint<float>("T"),
    AddBOp<GPUDevice, float>);

Я думаю, что регистрация AddB для GPU создает экземпляр объекта, который соответствует первой реализации, а не второй (первая реализация имеет два аргумента шаблона, вторая реализация имеет один).

вы, вероятно, могли бы исправить это, вызвав AddBOp во второй регистрации, хотя я бы посоветовал лучшие имена, чтобы избежать путаница.


Я только что прочитала TensorFlow issue #2054-ручное размещение на GPU пользовательского оператора с реализацией CPU и GPU всегда будет запускать версию CPU и поведение запуска реализации CPU, по-видимому, является особенностью TensorFlow, называемой "постоянным складыванием". Когда TensorFlow оптимизирует график перед первым запуском, операции с константами обычно оцениваются на CPU, поскольку считается, что реализации CPU и GPU должны давать одинаковые результаты. Обретать смысл.

два способа отключения этого поведения являются:

  1. отключение оптимизации графика:

    from __future__ import print_function
    import numpy as np
    import tensorflow as tf
    from time import time
    
    NUM_VALUES = 10
    
    input = np.arange(0, NUM_VALUES, dtype = float)
    
    custom_ops_module = tf.load_op_library('custom_ops.so')
    
    config = tf.ConfigProto(log_device_placement = True)
    config.graph_options.optimizer_options.opt_level = -1
    
    with tf.Session(config = config):
      start = time(); print(custom_ops_module.add_b(input, 8.).eval()); end = time(); print(end - start)
    
  2. не используя константы, например, путем подачи значений в заполнители:

    from __future__ import print_function
    import numpy as np
    import tensorflow as tf
    from time import time
    
    NUM_VALUES = 10
    
    custom_ops_module = tf.load_op_library('custom_ops.so')
    
    graph = tf.Graph()
    with graph.as_default():
      input = tf.placeholder(tf.float64, shape = (NUM_VALUES,))
      b = tf.placeholder(tf.float64, shape = ())
      result = custom_ops_module.add_b(input, b)
    
    with tf.Session(graph = graph, config = tf.ConfigProto(log_device_placement = True)) as session:
      feed_dict = {
        input: np.arange(0, NUM_VALUES, dtype = float),
        b: 8.,
      }
      start = time(); print(session.run([result], feed_dict = feed_dict)); end = time(); print(end - start)
    

по данным этой это может быть связано с управлением фрагментацией памяти, попробуйте:

with tf.device('/gpu:0'):

или фрагменты на связанной странице для настройки параметров фрагментации памяти.

EDIT: чтобы увидеть, если это так, попробуйте :

from __future__ import print_function
import numpy as np
import tensorflow as tf
from time import time

NUM_VALUES = 10

input = np.arange(0, NUM_VALUES, dtype = float)

custom_ops_module = tf.load_op_library('custom_ops.so')

config = tf.ConfigProto(log_device_placement = True)
config.gpu_options.allow_growth = True

with tf.Session(config = config):
    start = time(); print(custom_ops_module.add_b(input, 8.).eval()); end = time(); print(end - start)