Примечание редактора: Джос ван де Вольфшаар, инженер-исследователь из Королевского технологического института (KTH) в Швеции, недавно написал руководство, в котором показано, как реализовать пользовательские операции TensorFlow на основе CUDA, используя в качестве примера капсульные сети.
вводить
Недавно я немного поработал с программированием на CUDA и реализовал несколько пользовательских операций для TensorFlow. В качестве упражнения я решил попробовать настроить операцию в капсульной сети, которая потребовала бы некоторых хитростей с изменением формы или по крайней мере двух или трех промежуточных операций TensorFlow для завершения. Если вы не знакомы с капсульными сетями, вы можете прочитатьэта статья. Первоначальный автор статьи открыл исходный код капсульной сети. Код, обсуждаемый в этой статье: https://github.com/jostosh/capsnet
Многие из концепций, рассмотренных в этой статье (CUDA, пользовательские операции TensorFlow, тестирование градиента), можно изучить из соответствующей документации, но я всегда нахожу поучительным посмотреть, как эти отдельные элементы сочетаются друг с другом. Это также является основной мотивацией этой статьи: показать, как разработать пользовательскую операцию TensorFlow на основе CUDA от начала до конца. Итак, приступим.
Капсульное предсказание
Операции, рассматриваемые в этой статье, следующие:
Среди них Юj|iсостоит из капсулiПутем умножения матрицы на вектор WijuiКапсулы "Предсказание"jвектор активации. матрица Втijформа[out_dim, in_dim], вектор uiдля капсулiвыходной вектор . Один капсульный слой - все в одной партииi,jВыполните эту операцию на всех образцах. Следовательно, тензорW_ijиu_iФактическая форма[batch_size, in_caps, in_dim]и[in_caps, out_caps, out_dim, in_dim]. Это означает, что мы создадим операцию, которая берет эти два тензора, а затем вычисляет выходной тензор.u_hat_ji, форма этого выходного тензора[batch_size, in_caps, out_caps, out_dim]. Другими словами, нам нужно проиндексировать все пакеты.[0,1,...,batch_size-1]и все входные капсулы[0,1,...,in_caps-1]и выходной капсюль[0,1,...,out_caps-1]Комбинация вычисляет произведение матрицы на вектор.
Реализация ядра TensorFlow
Реализация пользовательских операций на основе графического процессора является наиболее ценной. Документация TensorFlow предоставляет необходимые материалы для написания ядер C++ для TensorFlow. После этого вы можете прочитать книгу CUDA на примере: введение в программирование графического процессора общего назначения или учебник Nvidia Developer's Even Easier Introduction to CUDA. В этой статье не будут повторяться эти детали, но будет представлен практический пример, который, надеюсь, поможет вам понять, как писать собственные операции TensorFlow на основе CUDA. Мне было легче изучить CUDA, чем я ожидал, но выжать всю производительность сложно, и я оставлю оптимизацию на потом.
Оп регистрации
Сначала давайте реализуем прямую связь Op (операции). Во-первых, мы зарегистрируем Op:
REGISTER_OP("CapsulePrediction")
.Input("input: T")
.Input("weights: T")
.Output("output: T")
.Attr("T: type")
.SetShapeFn([](InferenceContext* ctx) {
// 获取形状并确保维度正确
ShapeHandle in_shape;
ShapeHandle weights_shape;
TF_RETURN_IF_ERROR(ctx->WithRank(ctx->input(0), 3, &in_shape));
TF_RETURN_IF_ERROR(ctx->WithRank(ctx->input(1), 4, &weights_shape));
// 创建并设定输出形状
DimensionHandle out_d0, out_d1, out_d2, out_d3;
std::vector<DimensionHandle> out_dims;
out_dims.push_back(ctx->MakeDim(ctx->Dim(ctx->input(0), 0)));
out_dims.push_back(ctx->MakeDim(ctx->Dim(ctx->input(1), 0)));
out_dims.push_back(ctx->MakeDim(ctx->Dim(ctx->input(1), 1)));
out_dims.push_back(ctx->MakeDim(ctx->Dim(ctx->input(1), 2)));
ShapeHandle out_shape = ctx->MakeShape(out_dims);
ctx->set_output(0, out_shape);
return Status::OK();
});
скопировать кодВ настоящее время я определил их, и позже их можно добавить, добавив..: TТехнические характеристики разныеdtypeУкажите другое ядро TensorFlow.ShapeHandle,DimensionHandle,InferenceContextопределено вtensorflowПространства имен. Функция формы в коде реализована как лямбда-функция, которая сначала обеспечиваетctx->input(0)(введите васi)иctx->input(1)(вес Втij) имеет правильный ранг. Затем мы определяем размерность выходного тензора на основе входного тензора. Размер Op выхода[batch_size, in_caps, out_caps, out_dim], так что мы начинаем с uiПриобретение тензораbatch_size,in_caps, от ВтijПриобретение тензораout_capsиout_dim.
Предсказание капсулы с прямой связью
Теперь давайте посмотрим на Opядерный. «Ядро» в TensorFlow относится к реализации Op, соответствующей конкретному устройству. При определении пользовательского ядра оно должно начинаться с TensorFlow.OpKernelнаследовать и реализоватьComputeметод:
-
class CapsulePredictionOp : public OpKernel
-
{
-
public:
-
explicit CapsulePredictionOp(OpKernelConstruction* ctx) : OpKernel(ctx) { }
-
-
void Compute(OpKernelContext* ctx) override
-
{
-
// 获取输入
-
const Tensor& input = ctx->input(0);
-
const Tensor& weights = ctx->input(1);
-
-
// 设定输出形状
-
const TensorShape& input_shape(input.shape());
-
TensorShape output_shape(weights.shape());
-
output_shape.InsertDim(0, input_shape.dim_size(0));
-
output_shape.RemoveDim(4);
-
-
// 分配输出张量
-
Tensor* output = nullptr;
-
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, output_shape, &output));
-
-
// 获取本征（Eigen）张量并传给启动器
-
auto input_tensor = input.tensor<float, 3>();
-
auto weights_tensor = weights.tensor<float, 4>();
-
auto output_tensor = output->tensor<float, 4>();
-
launchCapsulePrediction(ctx->eigen_device(), input_tensor, weights_tensor,
-
output_tensor);
-
}
-
};
скопировать кодПриведенная выше реализация пока не затрагивает CUDA, но мы скоро это сделаем, так что не волнуйтесь. Приведенный выше код просто инициализирует выходную форму на основе входной формы и выделяет память. параметрOpKernelContextОбъект обеспечивает выделение памяти на используемом в данный момент устройстве. В нашем случае это будет GPU. Далее проходимtensorспособ получитьEigenтензоры и передать их нашим собственнымlaunchCapsulePredictionфункция (функция, выполняющая определенную работу).
запустить ядро
нашlaunchCapsulePredictionФункции буквально (по крайней мере, в терминах CUDA) запускают код на графическом процессоре. Возможно, это немного сбивает с толку, но CUDA относится к функциям, запускающим код на устройстве, какядерный. И в терминологии TensorFlow ядро не обязательно является реализацией графического процессора. Не будем слишком зацикливаться на этих терминах и сразу перейдем к коду:
void launchCapsulePrediction(
const GPUDevice& d,
typename TTypes<float, 3>::ConstTensor x,
typename TTypes<float, 4>::ConstTensor weights,
typename TTypes<float, 4>::Tensor out)
{
// 获取维度
const int64 batch_size = x.dimension(0);
const int64 in_caps = x.dimension(1);
const int64 in_dim = x.dimension(2);
const int64 out_dim = weights.dimension(2);
const int64 out_caps = weights.dimension(1);
// 维度一的尺寸
const int64 w_d0 = out_caps * out_dim * in_dim;
const int64 x_d0 = in_caps * in_dim;
const int64 o_d0 = in_caps * out_caps * out_dim;
// 维度二
const int64 w_d1 = out_dim * in_dim;
const int64 x_d1 = in_dim;
const int64 o_d1 = out_caps * out_dim;
// 维度三
const int64 w_d2 = in_dim;
const int64 o_d2 = out_dim;
// 为前馈操作运行CUDA核
CudaLaunchConfig config = GetCudaLaunchConfig(out.size(), d);
capsulePredictionKernel
<<<config.block_count, config.thread_per_block, 0, d.stream()>>>(
x.data(), weights.data(), out.data(),
o_d0, o_d1, o_d2, x_d0, x_d1, w_d0, w_d1, w_d2,
in_dim, out.size());
}
скопировать кодв параметрах функцииTTypesшаблоны иint64введитеtensorflowопределяется в пространстве имен. Следующий раздел, посвященный размерам, должен быть более или менее очевидным. Поскольку мы передаем тензорные данные на определенные ядра CUDA в виде одномерных массивов, нам нужно выяснить размер памяти для каждого измерения и каждого ядра. Обратите внимание, что когда я говорю «размер памяти», я имею в виду количество поплавков на ось, а не количество байтов. Рассмотрим размер памяти первой оси каждого тензора:
-
// 维度一的尺寸
-
const int64 w_d0 = out_caps * out_dim * in_dim;
-
const int64 x_d0 = in_caps * in_dim;
-
const int64 o_d0 = in_caps * out_caps * out_dim;
скопировать кодХорошо, так что мы можем просто использовать размеры, которые мы определили ранее. Код говорит нам,w_d0Но этоout_caps,out_dim,in_dimпродукт . Итак, если мы хотим начать с индекса Wi,j,k,lперейти к Wi+1,j,k,l, мы должныw_d0добавляется к одномерному индексу. Как вы, возможно, уже подумали, индексjиw_d1То же самое справедливо.
Код заканчивается исполнителем ядра CUDA:
// 为前馈操作运行CUDA核
CudaLaunchConfig config = GetCudaLaunchConfig(out.size(), d);
capsulePredictionKernel
<<<config.block_count, config.thread_per_block, 0, d.stream()>>>(
x.data(), weights.data(), out.data(),
o_d0, o_d1, o_d2, x_d0, x_d1, w_d0, w_d1, w_d2,
in_dim, out.size());
скопировать кодЭти утверждения включают в себя некоторые новые концепции. В первом утверждении используетсяGetCudaLaunchConfigпример для определенияблокироватьчисло, и каждый блокнитьколичество (см. файлы заголовков TensorFlowtensorflow/core/util/cuda_kernel_helper.h). Если вы планируете разработать свою собственную операцию, вам обязательно следует взглянуть на этот заголовочный файл.capsulePredictionKernelФункция распараллелена с помощью CUDA на графическом процессоре через тройной разделитель.<<<config.block_count, config.thread_per_block, 0, d.stream()>>>Бегать При запуске ядра необходимо указать количество блоков и количество потоков на блок. третий параметр0Пока это не имеет значения, если вы собираетесь реализовать собственное ядро, оно, скорее всего, будет иметь нулевое значение. CUDA-потокd.stream()Его можно рассматривать как конвейер инструкций графического процессора. Когда вы добавляете ядро в поток, поток гарантирует, что текущее ядро завершится до вызова следующего ядра. Если вы хотите запускать две независимые задачи параллельно, вы можете использовать два потока.
Потоки и блоки
Все назначены на один и тот же вызовкусокможно запускать параллельно. Если вы бежите сNблоки ядер, тогда вы можете думать об этом как о запущенномNнезависимые экземпляры функции ядра. Это так удобно!nvccКомпилятор гарантирует, что функция ядра обращается к точному индексу блока, чтобы конкретный экземпляр блока ядра знал, какие части массива, который он получает, обрабатываются.
Блок может содержать нескольконить. Потоки — это просто еще один уровень параллелизма, поэтому они могут работать параллельно. Вы спросите, а зачем придумывать еще один уровень? Потому что потоки могут делать то, что не могут блоки. Потоки могут совместно использовать память, что полезно, когда вы хотите использовать одно и то же значение из входного массива в одном и том же блоке несколько раз. Доступ к общей памяти намного быстрее и является одним из многих способов оптимизировать окончательную реализацию CUDA. Ниже приведена иллюстрация двух параллельных уровней:
Ядро CUDA
После краткого обзора потоков и блоков в CUDA мы, наконец, можем взглянуть на реализацию предсказания капсул с прямой связью в CUDA:
__global__ void capsulePredictionKernel(
const float* in, const float* weights, float* out,
const int64 o_d0, const int64 o_d1, const int64 o_d2,
const int64 x_d0, const int64 x_d1,
const int64 w_d0, const int64 w_d1, const int64 w_d2,
const int64 in_dim, const int64 output_size)
{
CUDA_1D_KERNEL_LOOP(i, output_size)
{
// 所以这里我们有 out[b,ci,cj,e]
const int64 b = i / o_d0;
const int64 ci = (i % o_d0) / o_d1;
const int64 cj = (i % o_d1) / o_d2;
const int64 e_out = i % o_d2;
// 接着我们可以看看如何计算in和W的数组索引
int64 in_idx = b * x_d0 + ci * x_d1;
int64 w_idx = ci * w_d0 + cj * w_d1 + e_out * w_d2;
// 初始化结果
float result = 0.0;
for (int64 v = 0; v < in_dim; ++v)
// 对`in`和`weights`而言，前馈计算的后续元素也处于内存中的后续位置
result += ldg(in + in_idx++) * ldg(weights + w_idx++);
// 写入结果
out[i] = result;
}
}
скопировать кодПервое, что вы заметите, вероятно, перед определением функции.__global__модификатор.nvccКомпилятор использует этот модификатор, чтобы сделать функциюГлобальныйДоступный означает, что он может быть запущен ЦП или хост-кодом в терминах CUDA. Параметры функции взяты изlaunchCapsulePredictionФункции наследуют имена, поэтому они не должны вызывать большой путаницы.CUDA_1D_KERNEL_LOOPдаtensorflow/core/util/cuda_kernel_helper.hОпределенный макрос:
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < output_size;
i += blockDim.x * gridDim.x)
скопировать кодЗапуск ядра CUDA с этим макросом из TensorFlow заставляет мыслить абстрактно и удобно: он дает нам некий индексi, индекс соответствует выходному массивуouti-й элемент . Различные экземпляры блоков/потоков ядра получат соответствующиеiсоответствующие им значения. Теперь нам просто нужно вычислить два других массиваinиweightsиндекс чего-либо. Чтобы вычислить этот индекс, мы определяем пакетный индексb, введите индекс капсулыci, индекс выходной капсулыcjи индекс элемента выходной капсулыe_out:
// 所以这里我们有 out[b,ci,cj,e]
const int64 b = i / o_d0;
const int64 ci = (i % o_d0) / o_d1;
const int64 cj = (i % o_d1) / o_d2;
const int64 e_out = i % o_d2;
скопировать кодОпределить их становится легко, когда мы знаем количество элементов, содержащихся на каждой оси. По сути, мы передаем объем памяти в качестве аргумента функции. Что касается других массивов, мы можем положитьb,ci,cj,e_outПреобразование в соответствующий одномерный индекс:
-
// 接着我们可以看看如何计算in和W的数组索引
-
int64 in_idx = b * x_d0 + ci * x_d1;
-
int64 w_idx = ci * w_d0 + cj * w_d1 + e_out * w_d2;
скопировать кодОпять же, мы получаем одномерный индекс, используя объем памяти, уже предоставленный каждой осью. Это (i) индексы партии, какbвходная капсулаciи (ii) входные капсулыci, выходная капсулаcj, элемент выходной капсулыe_outВес первого входного элемента капсулы. Если вы знакомы с Matlab, это в основном то, что есть в Matlab.sub2indфункция.
мы предполагаемinиWПоследняя ось соответствует входному элементу капсулы. Это означает, что они находятся в последующих местах в памяти, поэтому построить цикл, перебирающий элементы одной входной капсулы, несложно:
-
// 初始化结果
-
float result = 0.0;
-
for (int64 v = 0; v < in_dim; ++v)
-
// 对`in`和`weights`而言，前馈计算的后续元素也处于内存中的后续位置
-
result += ldg(in + in_idx++) * ldg(weights + w_idx++);
-
// 写入结果
-
out[i] = result;
скопировать кодldgФункция загружает функцию для кэша данных только для чтения. Он принимает указатель на элемент для чтения. Не забывайте, что мы вычисляем матрично-векторные произведения, которые представляют собой наборы внутренних произведений. Потенциальное улучшение заключается в использованииОбщая память, потому что значение одной входной капсулы используется много раз, но эту оптимизацию мы оставим на будущее.
контрольная работа
Я надеюсь, что эта статья является исчерпывающим примером разработки пользовательских операций (Ops) TensorFlow для графических процессоров. Это включает в себяТестовая операция. Следующее основано наnumpyОперация прямой связи выполняется:
-
import tensorflow as tf
-
from ops.capsuleprediction import capsule_prediction
-
import numpy as np
-
from parameterized import parameterized
-
import itertools
-
-
-
class CapsulePredictionOpTest(tf.test.TestCase):
-
-
@staticmethod
-
def _numpy_capsule_prediction(x, weights):
-
""" 使用numpy为x和weights生成输出 """
-
batch_size, in_caps, in_dim = x.shape
-
_, out_caps, out_dim, _ = weights.shape
-
-
out_shape = (batch_size, in_caps, out_caps, out_dim)
-
out = np.zeros(out_shape)
-
-
for b in range(batch_size):
-
for i in range(in_caps):
-
for j in range(out_caps):
-
for c in range(out_dim):
-
out[b, i, j, c] = np.dot(x[b, i], weights[i, j, c])
-
return out
скопировать кодops/capsuleprediction.pyВключаютcapsule_prediction(фактически отвечает за загрузку Op из скомпилированной общей библиотеки). Интерпретация приведенной выше функции должна быть простой: мы перебираем партию, входную капсулу, выходную капсулу, выходной элемент капсулы, а затем вычисляем скалярное произведение для каждой комбинации в выходных данных. Результат будет использован для проверки опережающего действия Op. Еще одна вещь, о которой стоит упомянуть, это то, что мы унаследовалиtf.test.TestCaseсвоего рода. Этот класс предоставляет некоторые вспомогательные функции для тестирования TensorFlow.
Протестируйте предсказание капсулы с прямой связью:
@parameterized.expand([
(batch_size, in_caps, out_caps, in_dim, out_dim) for
batch_size, in_caps, out_caps, in_dim, out_dim in
itertools.product([4, 8], [4, 8], [4, 8], [4, 8], [4, 8])
])
def test_capsule_prediction_op(self,
batch_size,
in_caps, out_caps, in_dim, out_dim):
""" 测试capsmatmul（胶囊矩阵乘法）操作 """
x = np.random.rand(batch_size, in_caps, in_dim)
weights = np.random.rand(in_caps, out_caps, out_dim, in_dim)
truth = self._numpy_capsule_prediction(x, weights)
with self.test_session() as sess:
x_ph = tf.placeholder(tf.float32, x.shape)
w_ph = tf.placeholder(tf.float32, weights.shape)
ret = capsule_prediction(x_ph, w_ph)
out = sess.run(ret, {x_ph: x, w_ph: weights})
self.assertAllClose(truth, out)
скопировать кодЗдесь используется множество трюков. Во-первых,parameterizedДекоратор предоставляет способ вызова тестов с разными параметрами, где каждый тест должен пройти успешно и будет вызванpytestСчитается независимым тестом. Если это не удается, данный ввод также отображается в журнале тестирования, поэтому, по моему опыту, использование этого декоратора ускоряет процесс отладки, если это необходимо.parameterized.expandОжидает список кортежей. Каждый кортеж будет расширен до позиционных аргументов функции. мы можем использоватьitertools.productЛегко сгенерировать множество кортежей с разными размерами,itertools.productПолучите декартово произведение всех параметров.
случайная инициализацияxиweightsмножество. Созданный нами график TensorFlow прост: он содержит только два заполнителя иcapsule_predictionСоч. Возвращаемое значение должно быть таким же, как_numpy_capsule_predictionзначение такое же. Запустим тест:
pytestПриятной особенностью является то, что вы можете использовать-kФлаг для выбора определенного набора тестов. Ура, все тесты пройдены!
Обратное предсказание капсулы
Далее обратная операция. Вы заметите, что некоторые концепции, с которыми мы столкнулись ниже, обсуждались ранее. Мы увидели, как вычислить правильный индекс одномерного массива по размеру измерения, написали ядро CUDA, зарегистрировали операцию и настроили тест. Так что с этого момента я собираюсь немного ускориться. То, что мы не обсудили, просто осталосьградиентточное определение. Сначала рассмотрим обычный плотный слой:
给定对应z的损失函数L的梯度，我们可以计算x和W的梯度：
где шi— вектор для строки i, а z' — вектор, хранящий локальные градиенты выходных данных. на основании вышеизложенногоiГрадиент , весь градиент x можно рассчитать по формуле:
Вычисление градиента операции прогнозирования капсулы такое же, нам просто нужно заменить z' на û'j|i, затем замените W на Wij.
Может быть проще вычислить градиенты для определенных весов:
Таким образом, матрица, хранящая градиент, является внешним произведением:
Интуитивно это говорит нам просто выбрать два соединенных нейрона из обоих слоев и умножить локальные градиенты выходного нейрона и входного нейрона. Это означает, что мы можем сделать то же самое для слоя прогнозирования капсул. Единственная разница заключается в размерности тензора по дизайну.
Градиент OpKernel
Не буду утомлять вас подробностями:
-
class CapsulePredictionGradOp : public OpKernel
-
{
-
public:
-
explicit CapsulePredictionGradOp(OpKernelConstruction* ctx) : OpKernel(ctx) { }
-
-
void Compute(OpKernelContext* ctx) override
-
{
-
// 获取输入张量
-
const Tensor& grad = ctx->input(0);
-
const Tensor& input = ctx->input(1);
-
const Tensor& weights = ctx->input(2);
-
-
// 获取形状以便分配输出
-
const TensorShape& input_shape(input.shape());
-
const TensorShape& weights_shape(weights.shape());
-
-
// 分配输出
-
Tensor* grad_input = nullptr;
-
Tensor* grad_weights = nullptr;
-
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, input_shape, &grad_input));
-
OP_REQUIRES_OK(ctx, ctx->allocate_output(1, weights_shape, &grad_weights));
-
-
// 获取本征（Eigen）张量并传给启动器
-
auto input_tensor = input.tensor<float, 3>();
-
auto weights_tensor = weights.tensor<float, 4>();
-
auto grad_tensor = grad.tensor<float, 4>();
-
auto grad_input_tensor = grad_input->tensor<float, 3>();
-
auto grad_weights_tensor = grad_weights->tensor<float, 4>();
-
launchCapsulePredictionGrad(
-
ctx->eigen_device<GPUDevice>(), input_tensor, weights_tensor, grad_tensor,
-
grad_input_tensor, grad_weights_tensor
-
);
-
}
-
};
скопировать кодЗдесь нет ничего нового. Важным отличием является то, что теперь нам нужно назначить два выходных тензора: один для весовых градиентов и один для входных градиентов. Форма градиента тензора такая же, как форма самого тензора. Таким образом, установка формы выделенных тензоров — это кусок пирога. Давайте проверим этоlaunchCapsulePredictionGradфункция:
void launchCapsulePredictionGrad(
const GPUDevice& d,
typename TTypes<float, 3>::ConstTensor input,
typename TTypes<float, 4>::ConstTensor weights,
typename TTypes<float, 4>::ConstTensor grad,
typename TTypes<float, 3>::Tensor grad_input,
typename TTypes<float, 4>::Tensor grad_weights)
{
const int64 batch_size = input.dimension(0);
const int64 in_caps = input.dimension(1);
const int64 in_dim = input.dimension(2);
const int64 out_dim = weights.dimension(2);
const int64 out_caps = weights.dimension(1);
// 维度一的尺寸
const int64 w_d0 = out_caps * out_dim * in_dim;
const int64 x_d0 = in_caps * in_dim;
const int64 o_d0 = in_caps * out_caps * out_dim;
// 维度二
const int64 w_d1 = out_dim * in_dim;
const int64 x_d1 = in_dim;
const int64 o_d1 = out_caps * out_dim;
// 维度三
const int64 w_d2 = in_dim;
const int64 o_d2 = out_dim;
// 运行输入梯度核
CudaLaunchConfig config = GetCudaLaunchConfig(grad_input.size(), d);
capsulePredictionInputGradKernel
<<<config.block_count, config.thread_per_block, 0, d.stream()>>>(
grad.data(), weights.data(), grad_input.data(),
w_d0, x_d0, x_d1, o_d0, o_d1, out_caps, out_dim, in_dim,
grad_input.size());
// 运行权重梯度核
config = GetCudaLaunchConfig(grad_weights.size(), d);
capsulePredictionWeightsGradKernel
<<<config.block_count, config.thread_per_block, 0, d.stream()>>>(
grad.data(), input.data(), grad_weights.data(), batch_size,
grad_weights.size(), w_d0, w_d1, w_d2, x_d0, x_d1, o_d0, o_d1, o_d2);
}
скопировать кодОпять же, мы видим похожую структуру кода. Мы получаем размеры, определяем размер памяти и, наконец, запускаем два ядра вместо одного.
Я рад, что вы все еще читаете! Теперь все становится немного сложнее. Взгляните на входной градиент ядра CUDA:
-
__global__ void capsulePredictionInputGradKernel(
-
const float* grad, const float* weights, float* grad_input,
-
const int64 w_d0,
-
const int64 x_d0, const int64 x_d1,
-
const int64 o_d0, const int64 o_d1,
-
const int64 out_caps,
-
const int64 out_dim,
-
const int64 in_dim,
-
const int64 output_size)
-
{
-
CUDA_1D_KERNEL_LOOP(i, output_size)
-
{
-
// 所以这里我们有 in_grad[b,ci,e]
-
const int64 b = i / x_d0;
-
const int64 ci = (i % x_d0) / x_d1;
-
const int64 e_in = i % x_d1;
-
-
// 接着我们可以看看如何计算in和W的数组索引
-
int64 w_idx = ci * w_d0 + e_in;
-
int64 grad_idx = b * o_d0 + ci * o_d1;
-
-
// 初始化结果
-
float result = 0.0;
-
// 迭代cj和e_out，我们已经具备其他索引
-
for (int cj = 0; cj < out_caps; ++cj)
-
{
-
for (int e_out = 0; e_out < out_dim; ++e_out)
-
{
-
// 增加grad_idx可以得到梯度的下一个元素
-
result += ldg(grad + grad_idx++) * ldg(weights + w_idx);
-
// 访问下一个输出胶囊元素，可以得到权重的下一个元素，
-
// 也就是说我们将w_idx加上in_dim
-
w_idx += in_dim;
-
}
-
}
-
// 写入结果
-
grad_input[i] = result;
-
}
-
}
скопировать кодЯ немного прокомментировал код, чтобы улучшить его читабельность. Подобно нашему предыдущему ядру CUDA, эта функция определяет индекс относительной оси:b,ci,e_in. Затем эти индексы используются для расчетаwиgrad(выходной градиент) 1D индекс. Один входной нейрон используется не только для произведений матрицы-вектора, но и для всех матриц предсказания.
Поэтому нам нужно два контура, один для выходной капсулы и один для отдельных элементов капсулы. Теперь нам нужно перейти к следующей выходной капсуле во внутреннем цикле, а не просто увеличивать индекс весов. Это означает, что нам нужно добавлять индекс к каждой итерации.in_dim.
Ядро CUDA с градиентом веса реализовано как:
-
__global__ void capsulePredictionWeightsGradKernel(
-
const float* grad, const float* input, float* grad_weights,
-
const int64 batch_size, const int64 output_size,
-
const int64 w_d0, const int64 w_d1, const int64 w_d2,
-
const int64 x_d0, const int64 x_d1,
-
const int64 o_d0, const int64 o_d1, const int64 o_d2
-
)
-
{
-
CUDA_1D_KERNEL_LOOP(i, output_size)
-
{
-
// 所以这里我们有 w[ci,cj,e_out,e_in]
-
const int64 ci = i / w_d0;
-
const int64 cj = (i % w_d0) / w_d1;
-
const int64 e_out = (i % w_d1) / w_d2;
-
const int64 e_in = i % w_d2;
-
-
// 接下来，我们可以看下如何为in和grad计算数组索引
-
int64 input_idx = ci * x_d1 + e_in; // (b == 0)
-
int64 grad_idx = ci * o_d1 + cj * o_d2 + e_out; // (b == 0)
-
-
// 初始化结果
-
float result = 0.0;
-
// 我们仅仅遍历b，因为我们已经有了其他索引
-
for (int64 b = 0; b < batch_size; b++)
-
{
-
result += ldg(grad + grad_idx) * ldg(input + input_idx);
-
// 跳至下个batch可以得到新元素
-
input_idx += x_d0;
-
grad_idx += o_d0;
-
}
-
grad_weights[i] = result;
-
}
-
}
скопировать кодЗдесь мы применяем тот же трюк: мы вычисляем индексы тензорных осей, мы вычисляем одномерные индексы, а затем перебираем нужные элементы, чтобы получить результат. Чего я не упомянул раньше, так это того, что нам нужно вычислить сумму градиентов каждого веса для всех образцов в партии. Учитывая размеры памяти входного тензора и выходного тензора для оси 0, это просто.
Прежде чем интегрировать нашу реализацию градиента в TensorFlow, нам нужно зарегистрировать градиент с помощью нашегоCapsulePredictionОп:
@ops.RegisterGradient("CapsulePrediction")
def _capsule_prediction_grad(op, grad):
""" 为胶囊预测操作计算梯度 """
return op_module.capsule_prediction_grad(grad, op.inputs[0], op.inputs[1])
скопировать кодТеперь мы можем напрямую использоватьtf.gradients, то граф вычисления градиента будет содержать наши операции с градиентом. чудесный!
тестовый градиент
Мы подошли к заключительному этапу: тестируем градиент. На самом деле это не так сложно, как кажется. TensorFlow поставляется с собственным инструментом тестирования градиента, который мы будем использовать здесь. мы даемCapsulePredictionOpTestКласс добавляет следующие методы:
@parameterized.expand([
(batch_size, in_caps, out_caps, in_dim, out_dim) for
batch_size, in_caps, out_caps, in_dim, out_dim in
itertools.product([4, 8], [4, 8], [4, 8], [4, 8], [4, 8])
])
def test_capsule_prediction_weights_grad(self, batch_size, in_caps, out_caps,
in_dim, out_dim):
""" 测试对应权重的输出梯度 """
x = np.random.rand(batch_size, in_caps, in_dim)
weights = np.random.rand(in_caps, out_caps, out_dim, in_dim)
out_shape = (batch_size, in_caps, out_caps, out_dim)
with self.test_session():
x_ph = tf.placeholder(tf.float32, x.shape)
w_ph = tf.placeholder(tf.float32, weights.shape)
fd = {x_ph: x, w_ph: weights}
caps_out = capsule_prediction(x_ph, w_ph)
grad_w = tf.test.compute_gradient(
w_ph, weights.shape, caps_out, out_shape, extra_feed_dict=fd
)
self.assertAllClose(grad_w[0], grad_w[1], atol=1e-3, rtol=1e-3)
@parameterized.expand([
(batch_size, in_caps, out_caps, in_dim, out_dim) for
batch_size, in_caps, out_caps, in_dim, out_dim in
itertools.product([4, 8], [4, 8], [4, 8], [4, 8], [4, 8])
])
def test_capsule_prediction_input_grad(self, batch_size, in_caps, out_caps,
in_dim, out_dim):
""" 测试对应x的输出梯度 """
x = np.random.rand(batch_size, in_caps, in_dim)
weights = np.random.rand(in_caps, out_caps, out_dim, in_dim)
out_shape = (batch_size, in_caps, out_caps, out_dim)
with self.test_session():
x_ph = tf.placeholder(tf.float32, x.shape)
w_ph = tf.placeholder(tf.float32, weights.shape)
fd = {x_ph: x, w_ph: weights}
caps_out = capsule_prediction(x_ph, w_ph)
grad_x = tf.test.compute_gradient(
x_ph, x.shape, caps_out, out_shape, extra_feed_dict=fd
)
self.assertAllClose(grad_x[0], grad_x[1], atol=1e-3, rtol=1e-3)
скопировать кодtf.test.compute_gradientФункция определяет «теоретический» и числовой градиент. Числовые градиенты вычисляют числовые градиенты путем разности, в то время как наши Op-зарегистрированные градиенты вычисляют теоретические градиенты. Они должны быть близки к равным, поэтому мы используем унаследованные отtf.test.TestCaseиз
assertAllClose
Метод утверждает, что они близки. Вот результат:
Да здравствует! Это сработало!
Работает по классификации MNIST
В моей предыдущей статье я обсуждал применение капсульных сетей к проблеме классификации MNIST. Мы можем вставить в кодcapsule_prediction:
-
def
digit_caps(incoming,
n_digit_caps,
dim_digit_caps,
name=
"DigitCaps"
,
-
neuron_axis=-
1
,
capsule_axis=-
2
,
routing_iters=
3
):
-
"""
数字胶囊层
"""
-
with
tf.variable_scope(name):
-
#
获取前一层的胶囊数和维度
-
in_shape
=
incoming.shape.as_list()
-
n_primary_caps
=
in_shape[capsule_axis]
-
dim_primary_caps
=
in_shape[neuron_axis]
-
#
初始化所有权重矩阵
-
if
args.custom_op:
-
w_shape
=
[
-
n_primary_caps,
n_digit_caps,
-
dim_digit_caps,
dim_primary_caps
-
]
-
else
:
-
w_shape
=
[
-
n_primary_caps,
-
n_digit_caps
*
dim_digit_caps,
-
dim_primary_caps
-
]
-
-
-
W_ij
=
tf.get_variable(
-
"weights"
,
shape=w_shape,
-
initializer=tf.keras.initializers.glorot_uniform()
-
)
-
#
初始化路由logit，
-
#
加上一个尺寸为1的第一轴，这样比较方便
-
b_ij
=
tf.get_variable(
-
"logits"
,
shape=[
1
,
n_primary_caps,
n_digit_caps],
-
initializer=tf.zeros_initializer(),
trainable=args.logits_trainable
-
)
-
if
args.custom_op:
-
#
定制操作
-
u_hat
=
capsule_prediction(incoming,
W_ij)
-
else
:
-
#
reshape和转置的奇技淫巧
-
u_i
=
tf.transpose(incoming,
(
1
,
2
,
0
))
-
u_hat
=
tf.matmul(W_ij,
u_i)
-
u_hat
=
tf.reshape(
-
tf.transpose(u_hat,
(
2
,
0
,
1
)),
-
(-
1
,
n_primary_caps,
n_digit_caps,
dim_digit_caps)
-
)
-
-
def
capsule_out(b_ij):
-
"""
给定logit
b_ij，计算该层的输出。
"""
-
c_ij
=
tf.nn.softmax(b_ij,
axis=
2
)
-
s_j
=
tf.reduce_sum(
-
tf.reshape(c_ij,
(-
1
,
n_primary_caps,
n_digit_caps,
1
))
*
u_hat,
-
axis=
1
-
)
-
v_j
=
squash(s_j)
-
return
v_j
-
-
def
routing_iteration(iter,
logits):
-
"""
-
给定一组logit，
-
基于论文中定义的路由计算新logit。
-
"""
-
v_j
=
capsule_out(logits)
-
a_ij
=
tf.reduce_sum(tf.expand_dims(v_j,
axis=
1
)
*
u_hat,
axis=
3
)
-
logits
=
tf.reshape(logits
+
a_ij,
(-
1
,
n_primary_caps,
n_digit_caps))
-
return
[iter
+
1
,
logits]
-
-
#
计算路由
-
i
=
tf.constant(
0
)
-
routing_result
=
tf.while_loop(
-
lambda
i,
logits:
tf.less(i,
routing_iters),
-
routing_iteration,
-
[i,
tf.tile(b_ij,
tf.stack([tf.shape(incoming)[
0
],
1
,
1
]))]
-
)
-
#
路由结果的第二个元素包含我们的最终logit
-
v_j
=
capsule_out(routing_result[
1
])
-
-
return
v_j
скопировать кодТак как насчет производительности? Что ж, получается, что обучение с кастомными операциями почему-то медленнее, чем обучение с решейпом и транспонированием. Как я упоминал ранее, этот код можно дополнительно оптимизировать. Возможно, это будет содержанием моей следующей статьи. В любом случае, спасибо, что дочитали до конца, предложения и отзывы приветствуются.
Оригинальный адрес: https://jostosh.github.io/posts/capscuda.html
Связанное чтение:CapsNet Primer Series One: интуитивное понимание капсульных сетей
CapsNet Primer Series 2: как работают капсулы
Третья серия записей CapsNet: алгоритм динамической маршрутизации между капсулами
CapsNet Начало работы, серия 4: Капсульная сетевая архитектура
Серия CapsNet Getting Started: реализация капсульных сетей на основе TensorFlow