В предыдущей статье я попытался описать, шаг за шагом, от краткого к более сложному, реализацию слоев свертки нейронной сети, которая достаточно хорошо работает на самых разных графических устройствах от разных поставщиков с разными архитектурами. Обоснованием для этого было то, что во многих приложениях, помимо достаточного уровня производительности, гибкость для быстрого развертывания новых решений на различных платформах является более ценной, чем получение немного большей скорости с помощью специализированных решений. Более того, распространение таких общих реализаций в общих базах кода глубокого обучения может вывести отрасль из ее нынешнего состояния крайней фрагментации — и безумия с проприетарными решениями, направленными на то, чтобы обойти конкурентов. А более широкое развертывание приведет к более быстрому внедрению и инновациям в этой области.

Продолжая это повествование, на этот раз я попытаюсь описать реализацию обратного распространения градиента свертки, которая хорошо работает на различных графических устройствах от разных поставщиков (включая Intel, NVidia, AMD, Qualcomm, ARM и Apple). Вариант использования, который мы здесь рассматриваем, — это точная настройка на периферийных устройствах, которая, как правило, выполняется во время простоя с низким энергопотреблением, а не в режиме полного обучения с нуля. Здесь также более ценно быстрое развертывание новых технологий, чем абсолютная скорость на каком-либо одном устройстве, поэтому необходимо искать единственное среднее решение. Но стоит отметить, что в итоге мы использовали это унифицированное решение для полного обучения и на больших картах NVidia/AMD. Хотя скорость обучения имеет значение (т. е. «достаточной» производительности для обучения не существует), мы обнаружили, что усилия по поддержке различных оптимизированных версий для полного обучения не стоят того. Это связано с тем, что мы тратили больше времени на сбор данных и эксперименты с сетевыми топологиями и гиперпараметрами, чем на ожидание завершения отдельных обучающих заданий.

Я снова попытаюсь описать пошагово оптимизацию ядра от лаконичной до продвинутой. Представленный здесь код написан на OpenCL, но его легко построчно перенести на iOS Metal. Далее предполагается, что читатель знаком с базовой терминологией графических вычислений, OpenCL и сверточных нейронных сетей.

Краткая реализация

Краткий код для обратного распространения входного градиента свертки показан во фрагменте ниже. В литературе вы можете увидеть утверждения о том, что градиентная обратная связь свертки сама по себе является просто сверткой на градиенте с транспонированными весами. Это связано с тем, что градиент по отношению к входным данным определяет, насколько чувствительны конечные выходные данные сети к входным данным слоя, учитывая, насколько мы обнаружили, что они чувствительны к выходным данным слоя. Чтобы лучше понять это, рассмотрим выходную привязку, обозначенную кружком на приведенной ниже иллюстрации свертки прямого прохода 3x3, и соответствующие верхние левые привязки во входных и весовых тензорах. Чтобы вычислить значение выходного якоря, в прямом проходе мы просто берем скалярное произведение входных значений и значений веса, начиная с якоря и двигаясь вниз и вправо по высоте и ширине.

Но теперь в обратном распространении, когда мы знаем, насколько чувствительна окончательная точность сети к небольшим изменениям нашей выходной привязки, нам нужно следовать линии в обратном порядке, чтобы определить чувствительность к небольшим изменениям нашей входной привязки — умножая чувствительность выходной привязки. по весу в верхнем левом углу ядра (цепное правило). Тем не менее, это дает нам только часть чувствительности, внесенную через этот выходной якорь. Тот же входной якорь также влияет на точность через другие выходные данные слоя. Чтобы получить общую чувствительность точности к нашему входному якорю, нам нужно суммировать вклад по всем маршрутам, через которые он может повлиять на конечный результат сети. Но если вы проследите, как скользящее окно свертки проходит через входные данные для получения выходных данных, вы заметите, что, несмотря на то, что остальные значения веса расположены вниз и справа от нашего якоря веса, остальные выходные точки слоя, через которые наш входной якорь влияет на окончательную точность, фактически вверх и влево от нашего выходного якоря. С одной стороны, это происходит потому, что здесь, в обратном распространении, мы сохраняем привязку входных данных слоя, в то время как при прямой свёртке мы сохраняем привязанными выходные данные слоя. И именно из-за этого преобразования ввода/веса по сравнению с прямым проходом мы можем рассматривать вычисление входного градиента как свертку выходного градиента с транспонированными весами ядра.

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

Например, обратите внимание, что краевые условия в backprop усложняются. В частности, здесь нам нужно иметь достаточно входного заполнения (независимо от того, было ли какое-либо дополнение на прямом проходе), чтобы его выходные пространственные размеры соответствовали входным размерам слоя. Это связано с тем, что свертки с размером ширины ядра, превышающим единицу, естественным образом уменьшают размер своих выходных функций (когда не применяется дополнение). Это означает, что бэкпроп таких сверток имеет больший выходной пространственный размер (ширину и высоту), чем их вход. В этом также легко убедиться при осмотре. В то время как углы и края входных данных для свертки 3x3 явно влияют на выходные данные, их градиент будет потерян в обратном распространении без дополнения. Как мы увидим в коде, это вводит некоторые нежелательные проверки условий, чтобы избежать обращений к буферу за пределами границ, что усложняет оптимизацию.

Но даже при этом отступы — это только часть проблемы, связанной с бэкпропом. В реальных фреймворках все становится еще сложнее. В частности, важной оптимизацией прямого прохода является объединение функции активации с выходными данными сверток. Это связано с тем, что проще записать значения свертки в память, просто повернуться и сразу же прочитать их обратно только для небольшой поэлементной операции. На самом деле, во время прохождения прямой свертки сделать это было настолько просто, что мы не заметили этого. Но это нужно решить здесь, в backprop, так как это немного мешает работе.

Как показано на рисунке ниже, градиент, возвращаемый сверточному слою, на выходе которого есть активация, является градиентом после активации. Поэтому его нужно умножить на производную от функции активации, чтобы получить градиент, который входит в фактическую свертку. Итак, теперь мы объединяем входные данные со сверткой, а не выходные. И проблема с этим заключается в том, что обычно существует большая избыточность в доступе к входному тензору свертки (из-за повторного использования по глубине вывода и пространственного перекрытия ядра), как показано с более толстыми ссылками, входящими в свертки. Слияние чего-либо на входе свертки обычно означает, что это тоже нужно будет делать с избыточностью. И что еще хуже, для вычисления производной нелинейной функции часто требуется исходное входное значение (обозначенное диагональной линией от результата прямой свертки), к которому теперь также нужно будет обращаться с избыточностью. Поэтому по своей природе лучше вместо этого объединить эту операцию с выходными данными backprop слоя впереди (т. е. переместить границу слоя в backprop от светло-голубой пунктирной линии к темно-синей пунктирной линии), где ее можно выполнить за один раз на Написать.

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

К счастью, стоимость вычисления производных активации часто намного меньше, чем стоимость доступа к значениям градиента в первую очередь (особенно в случае ReLU). И поэтому повторное применение этих производных может стоить накладных расходов, чтобы избежать потери более простых монолитных сплавленных слоев. Это то, что мы предполагаем в приведенном ниже коде, и, таким образом, просто применяем производную активации (df_act) при получении градиентов. Нам нужно будет уделить первоочередное внимание минимизации повторного доступа к входному градиенту, поскольку мы продолжим оптимизацию этого кода. Еще раз, чтобы уменьшить объем, мы сосредоточимся только на свертках полной глубины, которые используются в сетях типа VGG/Resnet. И хотя код предназначен для полных тензоров с плавающей запятой, общие концепции применимы и к другим форматам данных.

Чтобы избежать путаницы, в дальнейшем все ссылки на «вход» и «выход» сверточного слоя подразумеваются в отношении прямого прохода слоя (как и все ссылки на них в коде). И все ссылки на функции свертки (независимо от того, реализуют ли они прямой или обратный проход) относятся к самой функции. Далее следует краткий код ядра. Он очень похож на краткий код прямого прохода. Основное отличие состоит в том, что функция активации (точнее, ее производная) теперь появляется на входе функции свертки, а не на ее выходе. Кроме того, на выходе функции нет смещения или нелинейной функции активации.

__kernel void Conv_GradBackp(
        __global float* input_buf,
        __global float* weight_buf,
        __global float* out_buf,
        __global float* next_g_buf,
        __global float* g_buf) {
    const int half_pad = pad_out/2;
    int w_ = get_global_id(0) % in_width_;
    int h_ = get_global_id(0) / in_height_;
    int in_c = get_global_id(1);
    int batch = get_group_id(2);
    uint dest = ((batch*in_depth_+in_c)*in_width_*in_height_)+(h_*in_width_)+w_;
    float grad = 0;
    if (-)
        … //later optimized versions go here
    else
    {
        for (int out_c = 0; out_c < out_depth_; out_c++) {
            float input_buf_sub[$KERN_SIZE*$KERN_SIZE];
            float weight_buf_sub[$KERN_SIZE*$KERN_SIZE];
            for (int y_ = 0; y_ < kernel_size_; y_++) {
                int oH = h_-(kernel_size_-y_-1);
                for (int x_ = 0; x_ < kernel_size_; x_++) {
                    int iW = w_ — (kernel_size_-x_-1);
                    int o_idx = ((batch*out_depth_+out_c)*out_width_*out_height_+
                        (oH + half_pad) * out_width_ +
                        (iW + half_pad));
                    float err = next_g_buf[o_idx];
                    float df_act = df_activation(out_buf[o_idx]);
                    input_buf_sub[y_*$KERN_SIZE+x_] = err * df_act;
                    weight_buf_sub[y_*$KERN_SIZE+x_] = weight_buf[(int)(in_c*out_depth_*kernel_size_*kernel_size_+
                         out_c*kernel_size_*kernel_size_+
                         y_*kernel_size_+x_)];
                }
            }
            for (uint i = 0; i<$KERN_SIZE*$KERN_SIZE; i++) {
                grad = mad(input_buf_sub[i], weight_buf_sub[i], grad);
            }
        }
    }
    g_buf[dest] = grad;
}

В заголовке ядра input_buf, weight_buf и next_g_buf являются соответственно ссылками на буферы, содержащие входные данные и веса слоя, а также градиенты от предыдущего слоя. А g_buf — это собственный буфер градиента слоя, в который он записывает данные. Как и в реализации краткого прямого прохода, здесь также каждый рабочий элемент вычисляет одно значение (с той разницей, что здесь мы вычисляем градиенты, видимые на входе слоя). Глобальные индексы здесь используются немного иначе, чем в прямом проходе — с индексом 0, определяющим пространственное положение входного слоя (то есть по ширине и высоте), индексом 1, определяющим глубину входного слоя, и индексом 2, определяющим пакет. Снова предполагается, что буферы имеют формат расположения памяти NCHW (ширина, затем высота, образующие самые низкие размеры, за которыми следуют глубина канала и номер пакета). Мы избегаем использования тензорного сворачивания и общей локальной памяти — преимущества которых, как правило, очень чувствительны к архитектуре кэша.
Центральная часть кода циклически перебирает глубину вывода слоя и на каждой итерации загружает два частных массива; один для части входного градиента, соответствующего пространственному положению рабочего элемента и входной глубине, а другой — для весов на входной глубине рабочего элемента и выходной глубине итерации. Оба массива равны по размеру квадрату ширины ядра свертки. Затем он вычисляет скалярное произведение двух частных массивов. Результат этих скалярных произведений накапливается по позициям глубины вывода слоя (по самому внешнему циклу), чтобы получить окончательную сумму — градиент во входной позиции, соответствующей этому рабочему элементу.
На рисунке ниже показан размер входные и весовые тензоры для такой функции свертки, которая реализует обратное распространение слоя свертки с точки зрения слоя, и то, как диапазоны этих тензоров зависят от выходного индекса. Мы еще вернемся к этому рисунку, чтобы вспомнить, как каждая оптимизация меняет ситуацию. Сравнивая эту иллюстрацию с иллюстрацией для прямых сверток (из предыдущей статьи), сходство должно быть очевидным.

Как упоминалось ранее, обратное распространение свертки требует заполнения ввода, поскольку его ввод на самом деле меньше, чем вывод. И это вводит необходимость проверки границ, чтобы избежать доступа за пределы границ (в качестве альтернативы накладным расходам на фактическое заполнение нулями). Эти индексные проверки включены и выделены в приведенном ниже коде, размещенном после каждой итерации цикла ширины и высоты ядра. Кроме того, чтобы избежать необходимости инициализации нулями локальных массивов (необходимой из-за проверок границ), мы перемещаем скалярное произведение непосредственно в основной цикл.

if (-)
… //later optimized versions go here
else
{
    for (int out_c = 0; out_c < out_depth_; out_c++) {
        for (int y_ = 0; y_ < kernel_size_; y_++) {
            int oH = h_-y_;
            if ((oH % stride_out != 0)||(oH < 0)) continue;
            oH /= stride_out;
            if (oH >= out_height_-pad) continue;
            for (int x_ = 0; x_ < kernel_size_; x_++) {
                int iW = w_-x_;
                if ((iW % stride_out != 0)||(iW < 0)) continue;
                iW /= stride_out;
                if (iW >= out_width_-pad) continue;
                int o_idx = ((batch*out_depth_+out_c)* out_width_*out_height_+
                    (oH + half_pad) * out_width_ +
                    (iW + half_pad));
                float err = next_g_buf[o_idx];
                float df_act = df_activation (out_buf[o_idx]);
                float weight = weight_buf[(in_c*out_depth_* kernel_size_*kernel_size_ +
                    out_c * kernel_size_ * kernel_size_ +
                    y_ * kernel_size_ + x_)];
                grad += err * weight * df_act;
            }
        }
    }
}

Здесь есть еще одна проверка индекса, которая связана с шагом, а не с заполнением. В прямом проходе шаг реализуется просто путем введения умножения в вычисление начальной позиции входной пространственной области (сохранение последовательного доступа внутри области). Но в backprop шаг слоя превращается в расширение (см. здесь для хорошей иллюстрации того, почему это так). И с ними немного сложнее справиться, поскольку они вводят шаг между входными пространственными элементами внутри ядра. И мы увидим, что это усложняет использование векторизации ввода для оптимизации использования шины данных.

Еще одна важная вещь, на которую следует обратить внимание, касается вложенности циклов и этих пространственных проверок условий. Поскольку их результат не зависит от позиции глубины вывода, с точки зрения рабочего элемента, очевидно, лучше переместить цикл по глубине в самую внутреннюю позицию вложенности, чтобы избежать их повторного выполнения. Мы займемся этим, но пока воздержимся от этого, так как сначала мы хотим максимально уменьшить потери шины данных с помощью векторизации. И наличие глубины в самой нижней позиции вложенного цикла не поможет в этом, поскольку расположение в памяти как нашего веса, так и входных данных пространственно ориентировано.

Кроме того, в зависимости от размера ядра стоимость проверок пространственных условий может быть минимизирована за счет «запоминания». В частности, мы поднимаем расчеты проверки условий из цикла, чтобы они выполнялись один раз, а затем сохраняем их в небольшом локальном буфере (cond_pass[][]). Как показано ниже, логический результат для каждой позиции ядра можно быстро найти для каждой глубины. Большинство компиляторов не умеют делать это сами. Мы предполагаем, что эта оптимизация выполняется вручную, но для краткости опустим заполнение буфера запоминания в дальнейшем.

bool cond_pass[kernel_size][kernel_size];
for (int y_ = 0; y_ < kernel_size_; y_++) {
    for (int x_=0; x_<kernel_size_; x_++) cond_pass[y_][x_]=false;
    int oH = h_-y_;
    if ((oH % stride_out != 0)||(oH < 0)) continue;
    oH /= stride_out;
    if (oH >= out_height_-pad) continue;
    for (int x_ = 0; x_ < kernel_size_; x_++) {
        int iW = w_-x_;
        if ((iW % stride_out != 0)||(iW < 0)) continue;
        iW /= stride_out;
        if (iW >= out_width_-pad) continue;
        cond_pass[y_][x_]=true;
    }
}
for (int out_c = 0; out_c < out_depth_; out_c++) {
    for (int y_ = 0; y_ < kernel_size_; y_++) {
        for (int x_ = 0; x_ < kernel_size_; x_++) {
            if (!cond_pass[y_][x_]) continue;
            int o_idx = …

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

Векторизация весов

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

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

Сосредоточившись на свертке 3×3 (распространенной в VGG и Resnet), основной цикл выглядит так, как показано ниже. Примечание: оставление вариантов ядер в блоке if-then-else улучшает читабельность и не приводит к накладным расходам во время выполнения, когда условия основаны на параметрах JIT.

if (-)
 … //later optimized versions go here
else if ((kernel_size_ == 3)) {
    int size = kernel_size_ * kernel_size_;
        for (int out_c = 0; out_c < out_depth_; out_c++) { 
        // load input and weight for this sub area
        __global float* weight_ptr = &(weight_buf[in_c * out_depth_ * size + out_c * size]);
        struct float9 {
            float8 low8;
            float top1;
        };
        union {
            float s[9];
            struct float9 f9;
        } uweights;
        uweights.f9.low8 = vload8(0, weight_ptr);
        uweights.f9.top1 = weights_buf[in_c*out_depth_*size+ out_c*size+8];
        for (int y_ = 0; y_ < kernel_size_; y_++) {
            int oH = h_-y_;
            for (int x_ = 0; x_ < kernel_size_; x_++) {
                int iW = w_-x_;
                if (!cond_pass[y_][x_]) continue;
                int o_idx = ((batch*out_depth_+out_c)*out_width_*out_height_+
                     (oH/stride_out+half_pad)*out_width_ +
                     (iW/stride_out+half_pad));
                float err = next_g_buf[o_idx];
                float df_act = df_activation(out_buf[o_idx]);
                float inp = err * df_act;
                grad = mad(inp, uweights.s[y_*3+x_], grad); 
            }
        }
    }
}

Мы сняли нагрузку с гирь из петли. Его скалярные нагрузки 3*3=9 теперь выполняются через одну vload8 и одну скалярную загрузку. Это заметно повышает производительность практически со всеми распространенными графическими архитектурами и компиляторами по сравнению со скалярной версией, поскольку значительно снижает потери шины данных. Обратите внимание, что здесь мы определяем структуру float9 и используем ее в объединении с сохранением в ней load8 и скаляра, чтобы обеспечить последующий доступ итератора к ее элементам при выполнении сверток MAD. Это позволяет нам выполнять MAD внутри цикла загрузки ввода, чтобы избежать дорогостоящей нулевой инициализации входного массива. Обратите внимание, что один оставшийся скалярный доступ означает, что мы все еще не полностью избавились от потерь шины данных. Кроме того, для некоторых компиляторов может потребоваться ручное развертывание цикла MAD, чтобы избежать косвенного обращения к итератору при доступе к элементам структуры uweights. Мы не будем воспроизводить это здесь, так как код значительно расширяется.

На изображении ниже показано, как тензор весов отображается в памяти в формате IOHW (ширина фильтра, затем высота, формирующая самые низкие измерения, за которыми следуют выходные данные, а затем входная глубина). Перечисления измерений представляют их порядок в каждом, при этом длина измерения 1 макета памяти представляет собой ширину шины данных (или размер строки кэша). Когда значения весов извлекаются индивидуально как скаляры (и соответствующая сборка, сгенерированная компилятором, также имеет тенденцию делать это), только 1/8 этой ширины шины данных будет эффективно использоваться при каждом доступе — тратя впустую остальную часть шины. Но с векторизацией количество отдельных обращений сокращается до количества оставшихся скаляров (отмеченных «s») и красных точек, которые указывают на начало блока в строке кэша, который может быть выбран вместе. Обратите внимание, что один оставшийся скаляр на глубину сдвигает начало следующего 9x9 на единицу, что приводит к тому, что даже векторизованные обращения охватывают 2 строки кэша (снижается их эффективность, поскольку требуется доступ к 2 строкам кэша). Поэтому здесь есть еще много возможностей для улучшения. И в этом обратное продвижение отличается от прямого прохода. Но давайте сначала обратимся к эффективности доступа к входному тензору.

Векторизация входной ширины

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

Другое отличие состоит в том, что выборки входных значений не являются единообразными для SIMD, и, поскольку мы разместили дорожки SIMD в самом низком измерении рабочей группы (которое соответствует ширине входных данных), между дорожками SIMD существует неявная векторизация. Другими словами, при доступе к вводу по своей природе меньше потерь шины данных (по крайней мере, в том смысле, что используется материал, полученный по шине). На рисунке ниже показано, как два следующих соседних рабочих элемента по ширине SIMD (данные, к которым они обращаются, отмечены темно-синей и пурпурной диагональными линиями) добавляются к внутренней ширине выборки векторизованных данных. Это также показывает, как перекрытие их входных областей может обеспечить повторное использование выборки данных.

Таким образом, явная векторизация входных нагрузок по ширине может быть более эффективной, чем выборка их в виде отдельных скаляров, и дает все больше преимуществ по мере увеличения ширины ядра. Но здесь проверка условий заполнения и расширения усложняет ситуацию. На прямом проходе, где у нас не было таких проверок условий, в этот момент мы просто подняли входные нагрузки из внутреннего цикла и преобразовали их в выборки vload3. Но здесь, чтобы избежать потенциальных нарушений доступа к памяти, нам нужно убедиться, что никакая часть векторной области не выходит за границы тензора. Фрагмент кода ниже показывает, как можно изменить цикл загрузки ввода, чтобы извлечь из этого выгоду (сравните с тем же из прямого прохода). Он вводит нежелательную проверку условия в основном цикле, чтобы при необходимости вернуться к скалярному доступу (этот путь кода используется только для нешаговых сверток). Несмотря на накладные расходы, это дает значительный прирост производительности в теснорах с широким набором функций (верхние свертки Resnet/VGG), где большая часть пространственных областей свертки находится внутри границ.

for (int y_ = 0; y_ < kernel_size_; y_++)
 {
     int oH = h_-y_;
     if (oH >= out_height_-pad) continue;
     if (oH < 0) continue;
     if ((w_<out_width_-pad) && (w_>=(kernel_size_-1)))
     {
         int o_idx = ((batch*out_depth_+out_c)*out_width_*out_height_+
             (oH + half_pad) * out_width_ +
             ((w_-(kernel_size_-1)) + half_pad));
         __global float* nextgi = &(next_g_buf[o_idx]);
         __global float* nextacti = & output_buf[o_idx];
         const float3 nextg = vload3(0, nextgi);
         const float3 nextact = vload3(0, nextacti);
         float inpx = nextg.x * df_sigmod(nextact.x);
         float inpy = nextg.y * df_sigmod(nextact.y);
         float inpz = nextg.z * df_sigmod(nextact.z);
         grad = mad(inpx, uweights.s[y_*3+2], grad);
         grad = mad(inpy, uweights.s[y_*3+1], grad);
         grad = mad(inpz, uweights.s[y_*3+0], grad);
     }
     else for (int x_ = 0; x_ < kernel_size_; x_++)
     {
         if (!cond_pass[y_][x_]) continue;
         int o_idx = ((batch*out_depth_+out_c)*out_width_*out_height_+
               (oH + half_pad) * out_width_ +
               (iW + half_pad));
         float err = next_g[o_idx];
         float df_act = df_activation(output_buf[o_idx]);
         float inp = err * df_act;
         grad = mad(inp, uweights.s[y_ * 3 + x_], grad);
     }
 }

Таким образом, в отличие от загрузки весов, которую мы в основном конвертировали в одну большую векторную нагрузку по ширине и высоте, здесь мы можем делать только векторные нагрузки по ширине, оставляя их фрагментированными по высоте и вблизи пространственных границ. Однако более важным здесь является повышение эффективности за счет повторного использования данных на дорожках SIMD, а не за счет предотвращения потерь шины данных. Далее мы рассмотрим, как дальше развивать эту концепцию повторного использования данных.

Но сначала подведем итоги, возвращаясь к приведенной ранее иллюстрации сглаженного диапазона выходного тензора. До сих пор с помощью простой векторизации нам удалось сделать наш доступ к тензору весов в 4,5 раза более эффективным, а доступ к входному тензору максимально в 3 раза эффективнее ( в извилинах 3×3) — хотя и уменьшается в зависимости от размера объекта.

Объединение глубины вывода по слоям

Как мы видели ранее, при доступе к данным о весе оставалась заметная доля неэффективности. Несмотря на то, что мы обращаемся к его элементам с последовательных адресов, они фрагментированы по итерациям цикла. Мы можем исправить это, как мы сделали на прямом проходе, читая выровненные по кешу блоки, которые охватывают позиции глубины вывода, во временные буферы пинг-понга. В приведенном ниже коде показана реализация этого (для размера ядра 3x3). Выходная часть свертки с глубиной выполняется шагами по 16 — каждый раз извлекается 9 локальных переменных с плавающей запятой (которые могут содержаться в регистровом файле большинства графических архитектур). Проверки условий и входной путь векторизации обрезаны для краткости. Обратите внимание, что мы добавили в эту версию условие, согласно которому глубина вывода должна быть кратна 16 (что часто встречается в сетях типов VGG/Resnet), и оно будет удалено при JIT-компиляции слоев, которые проходят условие, поскольку мы предполагаем вывод глубина известна во время компиляции.

 if (-)
 … //later optimized versions go here
 else if ((kernel_size_ == 3) && (out_depth_%16==0))
 {
     int size = kernel_size_*kernel_size_;
     for (int out_c = 0; out_c < out_depth_; out_c+=16)
     {
         // load input and weight for this sub area
         float16 ping, pong;
         //3x3 at 16 different depths
         struct float9 { float8 low8; float top1; };
         union { float s[9]; struct float9 f9; } uweights[16];
         __global float* weight_ptr = &(weights_buf[in_c*out_depth_* size+out_c*size+0*16]);
         ping = vload16(0, weight_ptr);
         uweights[0].f9.low8 = ping.s01234567;
         uweights[0].f9.top1 = ping.s8;
         pong = vload16(1, weight_ptr);
         uweights[1].f9.low8 = (float8)(ping.s9abc, ping.sdef, pong.s0);
         uweights[1].f9.top1 = pong.s1;
         uweights[2].f9.low8 = pong.s23456789;
         uweights[2].f9.top1 = pong.sa;
         ping = vload16(2, weight_ptr);
         uweights[3].f9.low8 = (float8)(pong.sbcde, pong.sf, ping.s012);
         uweights[3].f9.top1 = ping.s3;
         uweights[4].f9.low8 = ping.s456789ab;
         uweights[4].f9.top1 = ping.sc;
         pong = vload16(3, weight_ptr);
         uweights[5].f9.low8 = (float8)(ping.sdef, pong.s012, pong.s34);
         uweights[5].f9.top1 = pong.s5;
         uweights[6].f9.low8 = pong.s6789abcd;
         uweights[6].f9.top1 = pong.se;
         ping = vload16(4, weight_ptr);
         uweights[7].f9.low8 = (float8)(pong.sf, ping.s0123, ping.s456);
         uweights[7].f9.top1 = ping.s7;
         uweights[8].f9.low8 = ping.s89abcdef;
         pong = vload16(5, weight_ptr);
         uweights[8].f9.top1 = pong.s0;
         uweights[9].f9.low8 = pong.s12345678;
         uweights[9].f9.top1 = pong.s9;
         ping = vload16(6, weight_ptr);
         uweights[10].f9.low8 = (float8)(pong.sabcd, pong.sef, ping.s01);
         uweights[10].f9.top1 = ping.s2;
         uweights[11].f9.low8 = ping.s3456789a;
         uweights[11].f9.top1 = ping.sb;
         pong = vload16(7, weight_ptr);
         uweights[12].f9.low8 = (float8)(ping.scdef, pong.s0123);
         uweights[12].f9.top1 = pong.s4;
         uweights[13].f9.low8 = pong.s56789abc;
         uweights[13].f9.top1 = pong.sd;
         ping = vload16(8, weight_ptr);
         uweights[14].f9.low8 = (float8)(pong.sef, ping.s0123, ping.s45);
         uweights[14].f9.top1 = ping.s6;
         uweights[15].f9.low8 = ping.s789abcde;
         uweights[15].f9.top1 = ping.sf;
         for (int y_ = 0; y_ < 3; y_++) {
             for (int x_ = 0; x_ < 3; x_++){
                 for (int out_c_in = 0; out_c_in < 16; out_c_in++) {
                     if (!cond_pass[y_][x_]) continue;
                     int o_idx = ((batch*out_depth_+out_c+out_c_in) 
                                      * out_width_ * out_height_ +
                                 ((h_-y_)/stride_out+half_pad)*  
                                      out_width_ +
                                 ((w_-x_)/stride_out+half_pad));
                     float err = next_g[o_idx];
                     float df_act = 
                                  df_activation(output_buf[o_idx]);
                     float inp = err * df_act;
                     grad = mad(inp, uweights[out_c_in].s[y_*3+x_], grad);
                 }
             }
         }
     }
 }

Таким образом, собирая веса для 16 глубин вывода с помощью 9 последовательных операций vload16, которые охватывают глубины вывода, мы подсказываем компилятору, как эффективно получать данные по шине. Это показано на изображении ниже, где последовательные точки в макете памяти указывают на данные, которые векторно выбираются вместе, а разные цвета указывают на разные позиции глубины в весовом тензоре 3x3. Обратите внимание, что, сделав цикл 16-кратной глубины (out_c_in) самым внутренним, мы также сократили количество повторных проверок границ при обходе высоты и ширины ядра. Нам нужно будет пересмотреть порядок этих внутренних циклов, когда мы будем двигаться вперед с оптимизацией.

Хотя мы сделали очень похожую оптимизацию в прямом проходе, там была существенная разница в том, что нам нужно было сначала объединить рабочие элементы по глубине вывода, что дало нам повторное использование входных данных для объединенных рабочих элементов. Но здесь выходная глубина слоя на самом деле является входной глубиной свертки (по которой выполняется итерация каждого рабочего элемента). Поскольку мы тренируем веса, нам нужно предположить то же тензорное свертывание, что и при прямом проходе. Мы могли бы изменить свертывание, чтобы сделать входную глубину меньшим измерением после ширины и высоты ядра (OIHW вместо текущего IOHW) в обоих случаях, но это повлияет на эффективность загрузки при прямом проходе. Таким образом, хотя мы объединили выборку по глубине вывода слоя с большой векторизованной выборкой, мы еще не добились повторного использования выборки данных между рабочими элементами. Таким образом, возвращаясь к сглаженной иллюстрации диапазона выходного тензора, улучшения, которых мы достигли на данный момент, выглядят так, как показано ниже.

Объединение глубины ввода слоя

Как поясняется в пошаговом руководстве по прямому проходу на иллюстрации ниже, для фиксированного входного пространственного положения мы можем повторно использовать извлеченные входные значения, комбинируя рабочие элементы по выходной глубине. Это имеет одновременное преимущество, позволяя нам выполнять эффективную выборку значений веса с помощью больших векторизованных загрузок, которые не тратят впустую ширину шины данных (при условии размещения данных веса IOHW). Но в функции свертки, которая реализует обратное распространение градиента того же слоя, входная и выходная глубины тензора весов меняются местами. Поэтому мы векторизовали выборку весов, комбинируя выборки по итерациям глубины, не объединяя рабочие элементы.

Но для большинства архитектур нам по-прежнему необходимо уменьшить пропускную способность выборки памяти, чтобы она не была узким местом и позволяла полностью использовать исполнительные блоки. И оставшийся архитектурно-независимый способ сделать это — объединить рабочие элементы по некоторому измерению, чтобы извлеченные данные можно было повторно использовать в них. Но мы должны быть осторожны, как мы это делаем, чтобы избежать давления на регистры, которое может привести к тому, что компилятор заполнит регистры в памяти и ухудшит производительность. Чтобы лучше передать эту проблему, на рисунке ниже красным показано, как мы изменили псевдокод прямой свертки, чтобы одновременно добиться повторного использования входных данных и эффективной выборки векторизованного веса. (примечание: на практике мы вытащили входные выборки из цикла ядра, что позволило нам развернуть и чередовать поэлементную векторизованную выборку весов).

На рисунке ниже показано, с «вводом» и «выводом» с точки зрения функции свертки (а не слоя), как мы до сих пор достигли векторизации выборки веса здесь в реализации обратного распространения.

Теперь, чтобы также добиться повторного использования ввода, нам нужно изменить код с зелеными дополнениями, показанными ниже. Хитрость здесь в том, что мы получаем за один раз значения веса, в 16x16 раз превышающие произведение высоты и ширины ядра. В различных архитектурах разумный размер регистрового файла (который определяет количество регистров, которые компилятор сможет использовать для хранения информации) для каждого рабочего элемента составляет от 128 до 256 записей. Таким образом, даже для сверток ядра 1x1 мы приблизимся к размеру регистрового файла или превысим его, если компилятор не сможет развернуть и распределить все самостоятельно. И это становится еще сложнее с большими размерами ядра. Чтобы обойти это, обычно требуется больше кормить компилятор с ложки, что приводит к расширению кода. Поэтому я прибегну к показу сжатого псевдокода для этих вариантов. Обратите внимание, что в качестве альтернативы мы могли бы уменьшить степень комбинированной выборки веса и повторного использования ввода до такой степени, что все помещается в регистровый файл. Но это происходит за счет пропускной способности памяти и ограничивает достижимую производительность.

Первое, на что следует обратить внимание, помогая компилятору с этим кодом, это то, что если мы разворачиваем внутренние циклы, нам действительно не нужно одновременно извлекать все 16 последовательных значений глубины. Это легче увидеть при рассмотрении размера ядра 1x1, где итераторы цикла ширины и высоты ядра исчезают. Как показано ниже, разворачивая внутренний цикл глубины с 16 входами, мы по-прежнему должны иметь возможность эффективно векторно выбирать значения весов, в то время как выполнение этого по частям уменьшает живой диапазон регистров.

В качестве конкретного примера, изменение кода загрузки веса пинг-понга с более раннего (для свертки 3x3) на что-то вроде приведенного ниже существенно сокращает реальный диапазон значений веса вдвое. Это высвобождает место в регистре для одновременной загрузки весов для большей глубины вывода для каждого рабочего элемента, и это позволяет реализовать то, что нам нужно; повторное использование дорогостоящих входных значений.

 ping = vload16(0, weight_ptr);
 uweights[0].f9.low8 = ping.s01234567;
 uweights[0].f9.top1 = ping.s8;
 pong = vload16(1, weight_ptr);
 uweights[1].f9.low8 = (float8)(ping.s9abc, ping.sdef, pong.s0);
 uweights[1].f9.top1 = pong.s1;
 uweights[2].f9.low8 = pong.s23456789;
 uweights[2].f9.top1 = pong.sa;
 ping = vload16(2, weight_ptr);
 uweights[3].f9.low8 = (float8)(pong.sbcde, pong.sf, ping.s012);
 uweights[3].f9.top1 = ping.s3;
 uweights[4].f9.low8 = ping.s456789ab;
 uweights[4].f9.top1 = ping.sc;
 pong = vload16(3, weight_ptr);
 uweights[5].f9.low8 = (float8)(ping.sdef, pong.s012, pong.s34);
 uweights[5].f9.top1 = pong.s5;
 uweights[6].f9.low8 = pong.s6789abcd;
 uweights[6].f9.top1 = pong.se;
 ping = vload16(4, weight_ptr);
 uweights[7].f9.low8 = (float8)(pong.sf, ping.s0123, ping.s456);
 uweights[7].f9.top1 = ping.s7;
 for (int y_ = 0; y_ < 3; y_++)
     for (int x_ = 0; x_ < 3; x_++)
         for (int out_c_in = 0; out_c_in < 8; out_c_in++) {
             if (!cond_pass[y_][x_]) continue;
             …
         }
 uweights[8].f9.low8 = ping.s89abcdef;
 pong = vload16(5, weight_ptr);
 uweights[8].f9.top1 = pong.s0;
 uweights[9].f9.low8 = pong.s12345678;
 uweights[9].f9.top1 = pong.s9;
 ping = vload16(6, weight_ptr);
 uweights[10].f9.low8 = (float8)(pong.sabcd, pong.sef, ping.s01);
 uweights[10].f9.top1 = ping.s2;
 uweights[11].f9.low8 = ping.s3456789a;
 uweights[11].f9.top1 = ping.sb;
 pong = vload16(7, weight_ptr);
 uweights[12].f9.low8 = (float8)(ping.scdef, pong.s0123);
 uweights[12].f9.top1 = pong.s4;
 uweights[13].f9.low8 = pong.s56789abc;
 uweights[13].f9.top1 = pong.sd;
 ping = vload16(8, weight_ptr);
 uweights[14].f9.low8 = (float8)(pong.sef, ping.s0123, ping.s45);
 uweights[14].f9.top1 = ping.s6;
 uweights[15].f9.low8 = ping.s789abcde;
 uweights[15].f9.top1 = ping.sf;
 for (int y_ = 0; y_ < 3; y_++)
     for (int x_ = 0; x_ < 3; x_++){
         for (int out_c_in = 8; out_c_in < 16; out_c_in++) {
             if (!cond_pass[y_][x_]) continue;
             …
         }

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

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

Объединение по ширине

Как и в прямом проходе, мы также можем комбинировать рабочие элементы по пространственной ширине (входная ширина слоя при обратном распространении). Хотя это более эффективно для слоев с широкими картами объектов, это приводит как к более широкой векторизованной выборке входных значений, так и, что более важно, к большему повторному использованию значений весов. Реализация и результат очень похожи на прямой проход, поэтому я не буду воспроизводить их здесь. Стоит упомянуть об одной основной сложности, связанной с проверками условий индекса (для заполнения и шага). Поскольку пространственные координаты этих комбинированных рабочих элементов не совпадают, нам необходимо вычислить их проверки индекса отдельно. И если для этого использовать мемоизацию, это увеличит необходимое хранилище регистров.

Спасибо за чтение. И дайте мне знать, что добавить в комментарии.