Устройство модуля PuzzleLib на примере линейного слоя

Введение

В этом туториале мы рассмотрим, как написать новый слой (модуль) для библиотеки PuzzleLib. Для этого мы подробно разберём, как устроен линейный модуль.

Обязательные методы

Родительский для всех модулей класс – Module. Для реализации нового слоя, необходимо наследовать его от Module (или его потомков) и реализовать следующие методы:

  1. __init__ – конструктор. В нём инициализируются внутренние переменные;
  2. updateData – метод, осуществляющий inference, т. е. проход данных вперёд;
  3. updateGrad – метод, осуществляющий вычисление градиента на данные (бэкпроп для данных);
  4. accGradParams – метод, осуществляющий вычисление градиента на параметры модуля (бэкпроп для параметров);
  5. dataShapeFrom – метод, которому можно передать форму входных данных и получить, какого размера тогда получатся выходные данные. Например, все модули активаций оставляют размер неизменным, а вот MaxPool со страйдом и без паддинга в несколько раз уменьшает каждую карту тензора;
  6. checkDataShape – метод, который валидирует размеры входных данных;
  7. gradShapeFrom – аналогичен методу dataShapeFrom, но нужен для вычисления размера градиента по размеру данных;
  8. checkGradShape – метод, который валидирует размеры градиента.

Углублённая теория

Чтобы лучше понять методы updateGrad и accGradParams, советуем изучить туториал Что лежит в основе - автоматическое дифференцирование.


__init__

Разберём, что происходит в конструкторе. Параметры конструктора описаны в документации модуля Linear.

    def __init__(self, insize, outsize, wscale=1.0, useBias=True, initscheme=None, name=None,
                 empty=False, transpose=False):

Сперва происходит вызов родительского конструктора:

        super().__init__(name)

Дальше – вызывается метод registerBlueprint, который необходим для того, чтобы модуль Linear поддерживал сериализацию через Blueprint:

        self.registerBlueprint(locals())

Запоминаем параметры слоя во внутренние переменные:

        self.transpose = transpose
        self.useBias = useBias

Если флаг empty был True, то больше ничего не делаем – так как не требуется заполнять слой начальными значениями:

        if empty:
            return

Дальше определяются размеры матрицы весов и вектора биасов. Если флаг transpose был True, то размеры будут обращённые:

        if transpose:
            wshape = (outsize, insize)
            bshape = (insize, )
        else:
            wshape = (insize, outsize)
            bshape = (outsize, )

Далее идёт инициализация матрицы весов случайными значениями, при этом учитывается, какая была указана initscheme. По умолчанию веса инициализируются по методу "xavier":

        self.W = None
        if initscheme == "none":
            self.setVar("W", Variable(gpuarray.empty(wshape, dtype=np.float32, allocator=memPool)))
        else:
            if initscheme == "xavier" or initscheme is None:
                if transpose:
                    nwscale = wscale / math.sqrt(outsize)
                else:
                    nwscale = wscale / math.sqrt(insize)
                W = np.random.uniform(-nwscale, nwscale, wshape).astype(np.float32)
            elif initscheme == "he":
                if transpose:
                    nwscale = wscale * math.sqrt(2.0 / outsize)
                else:
                    nwscale = wscale * math.sqrt(2.0 / insize)
                W = np.random.normal(0.0, nwscale, wshape).astype(np.float32)
            elif initscheme == "gaussian":
                nwscale = wscale
                W = np.random.normal(0.0, nwscale, wshape).astype(np.float32)
            elif initscheme == "uniform":
                nwscale = wscale
                W = np.random.uniform(-nwscale, nwscale, wshape).astype(np.float32)
            else:
                raise ValueError("Unsupported init scheme")

            self.setVar("W", Variable(gpuarray.to_gpu(W, allocator=memPool)))

Инициализация

Поддерживаются следующие схемы инициализации: "xavier", "he", "gaussian", "uniform", "none". В Интернете есть хорошие статьи об этом. Про инициализацию Xavier'а, например, см. здесь

Наконец, инициализируются биасы (смещения):

        self.b = None
        if useBias:
            self.setVar("b", Variable(gpuarray.zeros(bshape, dtype=np.float32, allocator=memPool)))

Резюме

В целом, всё, что происходит внутри конструктора – очень специфично, зависит от модуля. Но общая схема такая: записали внутренние переменные, инициализировали веса.


updateData

Реализует проход данных вперёд – inference.

В случае модуля Linear – это очень короткий метод:

    def updateData(self, data):
        if not self.transpose:
            self.data = CuBlas.mulMatrixOnMatrix(data, self.W)
        else:
            self.data = CuBlas.mulMatrixOnMatrix(data, self.W, transpB=True)

        if self.useBias:
            MatVec.addVecToMat(self.b, self.data, axis=1, out=self.data)

Здесь используется CuBlas – библиотека линейной алгебры на CUDA, поэтому тут всё очень просто. Перемножаются матрица входных данных и матрица весов, после чего к результату прибавляется вектор биасов. В случае реализации вашего модуля, здесь нужно проделать всю математику для inference.

Прототипирование

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


updateGrad

Как и с updateData, здесь используется CuBlas:

    def updateGrad(self, grad):
        if not self.transpose:
            self.grad = CuBlas.mulMatrixOnMatrix(grad, self.W, transpB=True)
        else:
            self.grad = CuBlas.mulMatrixOnMatrix(grad, self.W)

Вектор градиента умножается на транспонированную матрицу весов. Почему именно так - см. Что лежит в основе - автоматическое дифференцирование.


accGradParams

Осуществляет вычисление градиента на параметры модуля (бэкпроп для параметров). В случае линейного слоя – это в точности перемножение транспонированной матрицы входных данных на градиент справа:

    def accGradParams(self, grad, scale=1.0, momentum=0.0):
        if not self.transpose:
            CuBlas.mulMatrixOnMatrix(self.inData, grad, out=self.vars["W"].grad, transpA=True,
                                     alpha=scale, beta=momentum)
        else:
            CuBlas.mulMatrixOnMatrix(grad, self.inData, out=self.vars["W"].grad, transpA=True,
                                     alpha=scale, beta=momentum)

В случае реализации вашего модуля, необходимо разобраться, как вычислить градиент на веса слоя, и реализовать сперва на numpy, потом и на CUDA/OpenCL.

Параметры scale и momentum используются различными алгоритмами оптимизации. Обычный SGD не использует масштабирование градиента (scale) и не накапливает градиент от прошлых итераций с помощью momentum.

Соответственно, дальше, если необходимо, вычисляется градиент на биас. Это просто суммирование градиента по столбцам:

        if self.useBias:
            CuBlas.sumOnMatrix(grad, out=self.vars["b"].grad, alpha=scale, beta=momentum)

dataShapeFrom

Очень простой метод: он должен по размеру входных данных возвращать, какие данные получатся на выходе.

    def dataShapeFrom(self, shape):
        if not self.transpose:
            return shape[0], self.W.shape[1]
        else:
            return shape[0], self.W.shape[0]

checkDataShape

Тоже простой метод. Ему передаётся размер shape, и он проверяет, правильный ли это размер входных данных.

    def checkDataShape(self, shape):
        if len(shape) != 2:
            raise ValueError("Data must be 2d matrix")

        if not self.transpose:
            if shape[1] != self.W.shape[0]:
                raise ValueError("Expected %d data dimensions, %d were given" % (self.W.shape[0], shape[1]))
        else:
            if shape[1]!= self.W.shape[1]:
                raise ValueError("Expected %d data dimensions, %d were given" % (self.W.shape[1], shape[1]))

gradShapeFrom

Обратный метод к dataShapeFrom: по размеру выходных данных возвращает, какие данные получатся на входе при бэкпропе.

    def gradShapeFrom(self, shape):
        if not self.transpose:
            return shape[0], self.W.shape[0]
        else:
            return shape[0], self.W.shape[1]

checkGradShape

Аналогичен checkDataShape, проверяет правильность размера градиента.

    def checkGradShape(self, shape):
        if len(shape) != 2:
            raise ValueError("Grad must be 2d matrix")

        if not self.transpose:
            if shape[1] != self.W.shape[1]:
                raise ValueError("Expected %d grad dimensions, %d were given" % (self.W.shape[1], shape[1]))
        else:
            if shape[1] != self.W.shape[0]:
                raise ValueError("Expected %d grad dimensions, %d were given" % (self.W.shape[0], shape[1]))

Юнит-тесты

Жизненно необходимо писать юнит-тесты: в них вычисления должны быть реализованы на CPU, и результаты должны сравниваться с вычислениями на GPU. Желательно проверять вычисления на случайных тензорах случайных размеров.

Обязательно тестируем forward и backward вычисления. В Linear это осуществляется функциями calcTest и trainTest.

calcTest

Для проверки правильности вычислений вперёд, нам понадобятся "фейковые" обучающие данные: data и target (т. е. данные и лейблы). Генерируем их и отправляем на GPU:

def calcTest():
    insize = 5
    outsize = 1

    data = gpuarray.to_gpu(np.random.normal(0.0, 0.01, (5, insize)).astype(np.float32))
    target = gpuarray.to_gpu(np.random.normal(0.0, 0.01, (5, outsize)).astype(np.float32))

Затем строим линейный слой нужного размера и импортируем функцию ошибки, которая будет сравнивать результат вычислений с target:

    linear = Linear(insize, outsize)

    from PuzzleLib.Cost.MSE import MSE
    mse = MSE()

Пропускаем данные через него вперёд:

    linear(data)

Считаем ошибку и градиент:

    error, grad = mse(linear.data, target)

Делаем бэкпроп:

    linear.backward(grad)

Теперь делаем форвард и бэквард на CPU с помощью numpy:

    hostOutData = np.dot(data.get(), linear.W.get()) + linear.b.get()[np.newaxis, :]
    hostInGrad = np.dot(grad.get(), linear.W.get().T)

    hostWGrad = np.dot(data.get().T, grad.get())
    hostBGrad = np.sum(grad.get(), axis=0)

Сравниваем результаты GPU- и CPU-вычислений с помощью функции np.allclose, которая допускает малую погрешность вычислений:

    assert np.allclose(hostOutData, linear.data.get())
    assert np.allclose(hostInGrad, linear.grad.get())
    assert np.allclose(hostWGrad, linear.vars["W"].grad.get())
    assert np.allclose(hostBGrad, linear.vars["b"].grad.get())
Если бы что-то не сошлось, то тест упал бы в assert.

Тесты

В других, более сложных, модулях может потребоваться больше тестов с более разнообразными данными (см. свёрточный слой).

trainTest

В этом тесте проверяем, что бэкпроп работает.

def trainTest():
    insize = 500
    outsize = 100

    data = gpuarray.to_gpu(np.random.normal(0.0, 1.0, (32, insize)).astype(np.float32))
    target = gpuarray.to_gpu(np.random.normal(0.0, 1.0, (32, outsize)).astype(np.float32))

    linear = Linear(insize, outsize)

    from PuzzleLib.Cost.MSE import MSE
    mse = MSE()

    for i in range(100):
        learnRate = 1e-4

        linear(data)
        error, grad = mse(linear.data, target)

        linear.backward(grad)
        linear.updateParams(learnRate)

        if (i+1) % 5 == 0:
            print("Iteration #%d error: %s" % (i+1, error))

Заключение

В целом – всё. Изучайте, как реализуются другие, более сложные модули, и тогда вам будет проще реализовывать и свои новые модули.