ML на клиенте: Часть 1. Тренируемся на белочках.

Все, что вам нужно знать о цифровых белочках

В этой серии постов мы поговорим о ML-ной математике, а также о замечательных библиотеках, позволяющих эту математику применять. Чтобы немножко усложнить квест, я не буду использовать никаких формул и попробую объяснять все на белочках.

Я ни в коем случае не претендую на звание эксперта. Откровенно говоря, я совсем мало знаю про ML и DL :) Однако мои минимальные знания в этой области иногда помогают решать реальные продуктовые задачи. В общем – будем разбираться вместе :)

Начнем с вымышленной истории...

Маленький фронтенд который боялся математики

Жил-был фронтенд-разработчик. Он не сразу стал заниматься фронтом. Сначала он закончил крутые курсы, где узнал про React и Redux. Он устроился в крупную компанию, верстал там формочки и был счастлив. Но однажды к нему пришел грустный техлид проекта. Он рассказал, как в их чудесном сервисе для загрузки картинок белочек стали появляться изображения котиков и что это совершенно недопустимо. Нужен какой-то способ понять, что на картинке не белочка и предупредить пользователя об этой ужасной ошибке!

Фронтенд пошел гуглить. К сожалению, по запросу "js отличить белочку от котика" ничего толкового не нашлось. На Stackoverflow ему посоветовали использовать Tensorflow - крутой фреймворк для распознавания белочек. Фронтенд пошел читать документацию ... Сначала ему почему-то рекомендовали установить питон (а, как известно, некоторые фронтенды боятся змей). Потом фронтенд заглянул в следующий раздел документации и увидел вот такое:

Тензоры? Дата сеты? Слайсы? Фронтенд загрустил и решил делегировать задачу департаменту машинного обучения. Правда, потом он вспомнил, что такого департамента в их компании нет :)

A собственно, чего тут бояться?

Наш вымышленный фронтенд испугался математики :) Это нормально. Многие профессионалы в нашей области ее не любят.

Этот страх появляется еще в школе и развивается, превращаясь в уверенность "вот никогда-никогда никаких логарифмов и интегралов". На самом деле, бояться тут нечего. Математика значительно упрощают фронтожизнь.

Большинство читателей когда-нибудь пользовались CSS. Наверное, приходилось писать вот такое:

.squirrel {
  transform: translateX(100px);
}

Или вот такое:

.squirrel {
  transform: skewX(1);
}

А иногда и вот такое

.squirrel {
  transform: rotate(90deg);
}

На самом деле для всех случаев можно использовать один и тот же стиль:

.squirrel {
  transform: matrix3d(0, -1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, -100, 0, 0, 1);
}

Удобно – не надо думать, надо просто подставить нужные циферки в нужное место.

Давайте разберемся, что там такое написано.

Мatrix это двумерный массив:

;[
  [0, -1, 0, 0],
  [1, 0, 0, 0],
  [0, 0, 1, 0],
  [-100, 0, 0, 1],
]

И раз уж мы заговорили о массивах — массивы могут быть многомерными. В контексте ML мы называем многомерные массивы с элементами одного типа тензорами. Это не вполне точное определение, но нам нужно знать, что:

1. Тензор можно представить как многомерный массив
2. У тензора есть ранг и размерность. Ранг — это сколько индексов нам понадобится, чтобы достать элемент из тензора. Размерность — это количество элементов по каждой из осей. Наша матрица из примера — это тензор ранга 2 размерностью 3 x 3.
3. Как правило тензоры описывают преобразования между элементами какого-нибудь пространства.

Посмотрим, как это работает, на белочках.

Для каждого беличьего пискеля с координатами x y мы можем получить новые координаты, умножив x y на матрицу из примера выше.

Окей. С тензорами разобрались и математики больше не боимся.

Декларативно или императивно?

Есть два стиля написания кода: декларативный и императивный.

В первом случае у нас есть некая машинерия, которая позволяет формулировать проблемы на человеческом языке и описывать решения этих проблем в виде "я хочу, чтобы..."

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

Давайте сделаем что-нибудь императивное на старом-добром джаваскрипте. Напишем функцию, которая умеет складывать и перемножать циферки.

Смотрите, как императивно:

const createValueGetter = (a, b) => x => a * x + b
const getValue = createValueGetter(2, 5)
console.log([1, 2].map(getValue)) // [7, 9]

ML за 5 минут:

Мы можем действовать ровно наоборот. Создаем модель. Говорим ей, что у нас есть (inputs) и что мы хотим получить (realOutputs). Дальше мы как-то тренируем нашу модель. Натренированная модель будет выдавать нам нужный результат для любого инпута.

Смотрите, как декларативно:

const inputs = [1, 2, 3, 4]
const realOutputs = [7, 9, 11, 13] // values.map(getOutput)

// ... "скучный" код нашей нейросети

const learnedParams = doTrain()
const result = layer(inputs, ...learnedParams)

console.log(result) // [6.952457161700541, 8.9797075504127, 11.00695793912486, 13.034208327837018]

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

И тут у нас сразу же возникает пачка вопросов:

Что там внутри ящика?

Для этого давайте посмотрим на полную реализацию примера выше.

// императивненько
const createValueGetter = (a, b) => x => a * x + b
const getOutput = createValueGetter(2, 5)
console.log([1, 2].map(getOutput)) // [7, 9]

const inputs = [1, 2, 3, 4]

const realOutputs = [7, 9, 11, 13] // values.map(getOutput)
// const getValueWithML =  ??

// готовим декларативность :)

const trainStep = (a, b, inputs, realOutputs, step) => {
  const outputs = layer(inputs, a, b)
  const gradL = outputs.map((y, index) => y - realOutputs[index])

  const gradA = gradL.map((gr, i) => gr * outputs[i]).reduce((a, b) => a + b, 0)
  const gradB = gradL.reduce((a, b) => a + b, 0)

  return [a - gradA * step, b - gradB * step]
}

// задаем начальные параметры
const learningRate = 0.001
const numberOfSteps = 10000
const initialParams = [Math.random(), Math.random()]

// задаем нашу "архитектуру"
const layer = (inputs: number[], ...params: [number, number]): number[] =>
  inputs.map(x => params[0] * x + params[1])

// задаем как сравнивать результаты
const loss = (outputs: number[], realOutputs: number[]): number =>
  outputs
    .map((y, index) => Math.pow(y - realOutputs[index], 2))
    .reduce((a, b) => a + b, 0)

// 🏋️
const doTrain = (): [number, number] =>
  [...Array(numberOfSteps)].reduce((currentParams: [number, number]) => {
    console.log(
      ...currentParams,
      loss(layer(inputs, ...currentParams), realOutputs)
    )
    return trainStep(...currentParams, inputs, realOutputs, learningRate)
  }, initialParams)

const learnedParams = doTrain()

const result = layer(inputs, ...learnedParams)
console.log(result) // [6.952457161700541, 8.9797075504127, 11.00695793912486, 13.034208327837018]

TLDR:

Мы определяем "архитектуру ящика". Это набор функций, через которые мы будем прогонять наши входные данные. В нашем примере это:

const layer = (inputs: number[], a: number, b: number): number[] =>
  inputs.map(x => a * x + b)

Я не просто так назвала эту функцию слоем (layer).

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

const networkOutput =
  // ... layer 100
  layer3(layer2(layer1(inputs)))

У каждого слоя есть свои параметры. Наша задача их подобрать.

Здесь мы использовали очень простую функцию — слой. Реальные слои намного сложнее.

Как подобрать параметры?

Во-первых, нам нужно выдать сеточке метрику, которая будет показывать, насколько то, что она выдает, отличается от того, что мы хотим получить. Эта метрика называется функцией потерь (loss function).

В нашем случае это:

const loss = (outputs: number[], realOutputs: number[]): number =>
  outputs
    .map((y, index) => Math.pow(y - realOutputs[index], 2))
    .reduce((a, b) => a + b, 0)

Кстати, это что-то похожее на MSE

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

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

Мы учитываем градиенты при изменении параметров. Ведь чем больше градиент, тем сильнее соответствующий параметр влияет на результат.

Есть огромное количество ресурсов, где про это рассказывают подробно, например вот тут

И для каких случаев этот подбор будет успешным?

Есть прикольная теорема, называется универсальная теорема аппроксимации Она утверждает, что, используя определенный набор слоев, мы с нужной точностью можем приближать любую непрерывную функцию. То есть, если:

• у нас есть какая-то зависимость между входными и выходными параметрами;
• мы как-то можем посчитать потери;
• функция потерь достаточно симпатичная (простите меня, товарищи небезразличные к математике);

то наша сеть может научиться вести себя так же, как эта зависимость.

l1(x) = a0 * x + b0
l2(x) = a1 * x + b1

l2(l1(x)) = a1 * (a0 * x + b0) + b1
          = a1 * a0 * x + a1 * b0 + b1
          = a2 * x + b2

Это было просто. Давайте сделаем что-нибудь посложнее.

Как и в прошлый раз, начнем с императивного случая :)

Есть картинка белочки. Нужно ее заблюрить. Я не буду приводить здесь код преобразования пикселей в Uint8ClampedArray и рисования красивостей на canvas. Посмотрим только на блюрющую математику.

const kernel = [
  [1 / 16, 1 / 8, 1 / 16],
  [1 / 8, 1 / 4, 1 / 8],
  [1 / 16, 1 / 8, 1 / 16],
]

const convStep = (arr1: number[], kernel): number =>
  kernel.flat().reduce((acc, v, i) => acc + v * arr1[i], 0)

const convolve = (
  array: Uint8ClampedArray,
  kernel: number[][],
  w: number,
  h: number,
  stride = 1,
  chInImage = 4
): Uint8ClampedArray => {
  const result = new Uint8ClampedArray(w * h * chInImage).fill(255)
  const kh = kernel.length
  const kw = kernel[0].length

  for (let i = 0; i < w - kw; i += stride) {
    for (let j = 0; j < h - kh; j += stride) {
      for (let c = 0; c < chInImage; c++) {
        const arrToConsolve: number[] = []
        for (let k = 0; k < kw; k++) {
          for (let l = 0; l < kh; l++) {
            arrToConsolve.push(
              array[
                chInImage * w * j +
                  chInImage * i +
                  c +
                  chInImage * k +
                  chInImage * l * kw
              ]
            )
          }
        }

        const convStepResult = convStep(arrToConsolve, kernel)
        result[chInImage * w * j + chInImage * i + c] = convStepResult
      }
    }
  }

  return result
}

const blurredSquirrel = convolve(
  imagePixels,
  kernel,
  originalSquirrel.width,
  originalSquirrel.height
)

Мало того, что мы заблюрили белочку, мы еще и побили рекорд — 5 вложенных циклов for!

Все очень императивно, правда же? У нас есть готовые числа (их еще называют ядром (и мы теперь знаем, что это тензор :)). Мы берем эти числа и выполняем заранее известные преобразования. Для блюра белочки мы использовали операцию convolve. Эта операция берет матрицу заданного размера, выбирает кусочек картинки такого же размера и "сворачивает" их вместе.

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

Выбирая правильные ядра, мы можем по-разному преобразовывать нашу картинку:

• детектировать грани разных направлений,
• сглаживать или добавлять резкости,
• менять цвета, яркость и контраст.

Цифровые белки опасны!

Давайте посмотрим на другую белочку.

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

Если размер белочки M * N (2388 x 1668) и белочка черно-белая (мы используем только один канал), то нам понадобится матрица размером 2388 x 1668 x 1.

Работать с такими большими объемами данных тяжело. Для примера с блюром, время заблюривания линейно зависит от M x N Хорошие новости – нам не нужны все пиксели нашего изображения. Вместо этого мы можем представить нашу белочку как-то иначе.

Как выбрать правильное представление?

Если мы "расплющим" беличью матрицу, получится боооольшой массив длинной 3983184 (это примерно 3.8 Mb). Интересно, сколько пикселей из этих 3983184 нам понадобятся, чтобы уверенно отличать белку от других зверушек? Давайте вспомним, что никто никогда не хранит "сырые" пиксели и все картинки сжаты каким-нибудь добрым алгоритмом.

Например, если мы сохраним белочку в фромате SVG, мы получим всего 16Kb. При этом мы можем смотреть на наше изображение как на набор геометрических примитивов. Простых – кружочки и квадратики — которые в свою очередь формируют более сложные – беличий хвост.

Возможно, мы можем представить наши пиксели как-то иначе? Сделать так, чтобы одни пиксели были важнее других для нашей задачи? Получив такое представление, мы можем просто выкинуть все ненужное. Часть информации о белочке при этом потеряется, но при должной сноровке мы сможем восстановить белочку.

Давайте попробуем это сделать, использовав PCA Основная идея PCA - изменить представление наших данных так, чтобы их легче было разделить.

Например, для набора из 800 картинок нам будет достаточно 340 значений (это примерно 0.33 КБ), чтобы уверенно отличать нашу белочку от всего остального.

Вот так это выглядит. Слева - оригинал, справа - восстановленная из 340 циферок белка.

Зачем мы занимаемся этими странными штуками?

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

Настало время решить ту неподъемную задачу и наконец-то отличить белку от кота!

Итак, наш вымышленный фронтенд-разработчик сидит и грустит. Давайте ему поможем.

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

Первым делом нам нужно загрузить готовую модель. Вымышленно-нагугленное решение со stackoverflow нам тут как раз пригодится. Только мы будем использовать не питонячью, а джаваскриптовую версию tensorflow. Зовется она (вы не поверите!) tensorflowjs. Пока нам не важно, что это и как оно работает. Главное, что может модельки грузить :)

const loadModel = async (): Promise<tf.LayersModel> => {
  const model = await tf.loadLayersModel(MODEL_URL)
  return model
}

Отлично! Моделька загружена. Она умеет предсказывать вероятности нахождения на картинке определенных объектов. Эта модель тренировалась на огромном дата-сете ImageNet, состоящем из 1000 классов объектов. Все классы заботливо сложены в json и грузятся отдельно (IMAGENET_CLASSES это как раз оно).

Теперь давайте посмотрим, что у MobileNet внутри:

model.summary()

Ага! Слои присутствуют, параметры присутствуют и там еще наши старые-добрые функции convolve есть. Чтобы запихнуть белку, кота или кого-нам-еще-захочется в модель, надо написать хелпер, который после небольшого препроцессинга входной картинки выдаст нам нужные вероятности:

const MODEL_URL = '/models/mobilenet/model.json'
const PREPROCESS_DIVISOR = tf.scalar(255 / 2)
const WIDTH = 224
const HEIGHT = 224

// функция препроцессинга
const processInputImage = (input: tf.Tensor) => {
  const preprocessedInput = tf.div(
    tf.sub(input.asType('float32'), PREPROCESS_DIVISOR),
    PREPROCESS_DIVISOR
  )
  return preprocessedInput.reshape([-1, ...preprocessedInput.shape])
}

const predict = async (input: tf.Tensor, model: tf.LayersModel) =>
  model.predict(processInputImage(input))

const getProbs = async (
  img: HTMLImageElement,
  model: tf.LayersModel
): Promise<Predictions> => {
  // превращаем картинку в тензор
  const tensor = tf.browser.fromPixels(img)
  // предсказываем
  const result = await predict(tensor, model)
  const predictedClasses = tf.tidy(() => {
    // получаем самые вероятные предсказания
    const { values, indices } = tf.topk(result as tf.Tensor)
    const valuesData = values.dataSync() as Float32Array
    const indexData = indices.dataSync()
    return valuesData.reduce(
      (acc, val, idx) => [
        ...acc,
        {
          prob: val,
          // и превращаем циферку соответствующую классу в его название
          cl: IMAGENET_CLASSES[indexData[idx]][1],
        },
      ],
      []
    )
  })

  return predictedClasses
}

И вот что получается:

А что там за слои такие?

Давайте дальше разбираться в нашей модели. Как вообще она учится и что делает? Мы можем на это посмотреть :)

Для начала нам нужно посмотреть, какие вообще слои есть:

model.summary()

Мы можем вытащить каждый слой и поизучать его отдельно

const LAYER_NAME = 'block_12_expand'
const layer = model.getLayer(LAYER_NAME)

// или вот так
const layer100 = model.layers[100]

Как я упоминала выше, каждый слой — это просто функция. Мы можем ее спокойно применять.

const layer = model.layers[0]
const result = layer.apply(processInputImage(input))

В результате мы получим обычный тензор размерностью 1 x 112 x 112 x 32 и рангом 4. Но есть еще один интересный способ использования функции apply Мы можем вызвать эту функцию не на настоящем тензоре с готовыми циферками, а на "заглушке". Такая заглушка называется SymbolicTensor.

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

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

Сначала нам нужно выбрать интересующий нас слой

const LAYER_NAME = 'block_12_expand'
const layer = model.getLayer(LAYER_NAME)

Дальше разделить модельку на две. Первая — все слои до нужного нам, вторая - все слои после нужного нам и до конца. Вторую модель мы собираем через apply.

// находим нужный слой и его output
const layerOutput = layer.output
let layerIndex = model.layers.findIndex(l => l.name === LAYER_NAME)
const [_, ...outputShape] = (layerOutput as tf.SymbolicTensor).shape

// первая модель
const m1 = tf.model({
  inputs: model.inputs,
  outputs: layerOutput,
})

// вторая модель
const m2Input = tf.input({ shape: outputShape })

let nextTensor = newModelInput
const m2Layers = model.layers.slice(layerIndex + 1)
for (const l of m2Layers) {
  nextTensor = l.apply(nextTensor) as tf.SymbolicTensor
}
const m2 = tf.model({ inputs: newModelInput, outputs: nextTensor })

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

const CLASS_INDEX = 335
// функция, для которой будем считать градиенты. Возвращает вероятность получить интересующий нас класс для второй модели.
const classProbability = (input: tf.Tensor) =>
  (m2.apply(input, { training: true }) as tf.Tensor).gather([CLASS_INDEX], 1)

// собственно градиент
const gradFn = tf.grad(classProbability)

// прогоняем первую модель
const m1Output = m1.apply(input)

// считаем, как output интересующего нас слоя влияет на вероятность получить нужный класс
const gradValues = tf.mean(gradFn(m1Output as tf.Tensor), [0, 1, 2])

// применяем градиенты
const m2ScaledOutput = (m1Output as tf.Tensor).mul(gradValues)

// на основе градиентов строим тепловую карту
heatMap = getHeatMap(scaledConvOutputValues)

// ресайзим heat map
tf.image.resizeBilinear(heatMap as tf.Tensor<tf.Rank.R3>, [width, height])

Карта подсвечивает те части изображения, благодаря которым наш слой выдал такие значения, что получился нужный класс.

Фууух! Я немного утомилась, но продолжение обязательно последует

Кажется, наш вымышленный фронтенд доволен и его не уволят с работы. Более того, мы теперь чуть-чуть понимаем, как MobileNet и похожие сети распознают изображения. Но мы все еще очень далеки от понимания того, как применять эти знания и техники на клиенте. Так что продолжение следует!