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

Чтобы научиться писать пользовательские операции TensorFlow, я следовал Добавление нового оп учебник и сделал опцию «add_b», которая добавляет скаляр 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.

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

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 tenorflow / stream_executor / cuda / cuda_gpu_executor.cc: 892] OS X не поддерживает NUMA - возвращая ноль узла NUMA
I tenorflow / core / common_runtime / gpu / gpu_device.cc: 951] Найдено устройство 0 со свойствами:
название: GeForce GT 750M
мажор: 3 минор: 0 memoryClockRate (ГГц) 0,9255
pciBusID 0000: 01: 00.0
Общая память: 2,00 ГБ
Свободная память: 1,80 ГБ
I tenorflow / core / common_runtime / gpu / gpu_device.cc: 972] DMA: 0
I tenorflow / core / common_runtime / gpu / gpu_device.cc: 982] 0: Y
I tenorflow / core / common_runtime / gpu / gpu_device.cc: 1041] Создание устройства TensorFlow (/ gpu: 0) -> (устройство: 0, имя: GeForce GT 750M, идентификатор шины pci: 0000: 01: 00.0)
Отображение устройства:
/ job: localhost / replica: 0 / task: 0 / gpu: 0 -> device: 0, имя: GeForce GT 750M, идентификатор шины pci: 0000: 01: 00.0
I tenorflow / core / common_runtime / direct_session.cc: 252] Отображение устройства:
/ job: localhost / replica: 0 / task: 0 / gpu: 0 -> device: 0, имя: GeForce GT 750M, идентификатор шины pci: 0000: 01: 00.0

AddB: / job: localhost / replica: 0 / task: 0 / gpu: 0
Я тензорный поток / core / common_runtime / simple_placer.cc: 819] AddB: / job: localhost / реплика: 0 / task: 0 / gpu: 0
AddB / b: / job: localhost / replica: 0 / task: 0 / gpu: 0
Я тензор потока / 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 tenorflow / core / common_runtime / simple_placer.cc: 819] AddB / input: / job: localhost / replica: 0 / task: 0 / gpu: 0
[6. 5. 4. 3. 2.]

«Журналы размещения устройства», по-видимому, указывают на то, что операция выполняется на графическом процессоре, но вывод указывает, что используется реализация ЦП.

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

E tenorflow / core / common_runtime / executor.cc: 334] Исполнителю не удалось создать ядро. Не найдено: Нет зарегистрированного "AddB" OpKernel для устройств ЦП, совместимых с узлом AddB = AddB [T = DT_FLOAT, _device = "/ job: localhost / replica: 0 / task: 0 / gpu: 0"] (AddB / input, AddB / б)
, Зарегистрировано: устройство = 'GPU'; T в [DT_FLOAT]
Устройство = «ГПУ»; T в [DT_DOUBLE]

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

Это сообщение об ошибке выглядит для меня как ошибка, потому что, хотя в сообщении говорится, что «Исполнителю не удалось создать ядро», ядро, очевидно, было создано для запуска операции на графическом процессоре.

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

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

  • Я использую MacBook Pro со встроенным NVIDIA GeForce GT 750M (CUDA Compute Capability 3.0).
  • macOS Sierra Версия 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 или меньше, то используется реализация ЦП. когда NUM_VALUES 1310721 или больше, то используется реализация графического процессора.

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

8

Решение

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

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 < float> во второй регистрации, хотя я бы посоветовал лучшие имена, чтобы избежать путаницы.

1

Другие решения

Я просто читал Выпуск TensorFlow № 2054 — Ручное размещение на графическом процессоре пользовательского оператора с реализацией как CPU, так и GPU всегда будет запускать версию 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)
    
1

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

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

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

РЕДАКТИРОВАТЬ: Чтобы увидеть, если это так, попробуйте:

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