В этом посте надеюсь поделиться с вами некоторыми знаниями, которые усвоил при разработке этого проекта. Он заставил меня глубоко задуматься, как именно PyTorch работает под капотом. Возможно, вы для начала хотите уточнить, почему вообще важно понимать механику PyTorch. Возьмусь утверждать, что, не понимая этих абстракций или просто принимая их за чистую монету, вы можете занести в ваш код целую кучу багов. После этого он в лучшем случае окажется плохо оптимизирован, а в худшем нарушится ваш вычислительный граф, поэтому ваша модель никогда ничему не сможет научиться
Предполагаю, что вы хотя бы в общих чертах знакомы с обратным распространением, библиотекой PyTorch, а также понимаете, как работают компьютеры. Если чего-то из этого не хватает, то переходите по ссылкам, которые я буду оставлять в тексте. На их основе вы сможете интуитивно понять, как связаны эти концепции.
Что же такое тензор?
В контексте машинного обучения тензор — это n-мерная матрица. Хорошо, а что же тогда такое torch.Tensor? Точнее, что именно происходит, когда мы выполняем такой фрагмент кода: a = torch.tensor(1.0, requires_grad=True). Оказывается, что PyTorch выделит данные в куче и вернёт указатель на эти данные, причём, это уже будет разделяемый указатель. Чтобы лучше понять, как устроены указатели в PyTorch, рассмотрим несколько примеров.
a = torch.arange(9).reshape(3, 3) # 3x3
b = a.t() # Транспонирование
b[0, 0] = 9999
a
# >> tensor([[9999, 1, 2], [3, 4, 5], [6, 7, 8]])
Обратите внимание: в вышеприведённом примере .t() возвращает указатель, и это означает, что и a, и b направлены на один и тот же базовый фрагмент данных. Если как-то изменить эти данные, то внесённые изменения будут видны с обоих этих указателей.
a = torch.arange(9).reshape(3, 3) # torch.int64
b = a.to(torch.float16)
b[0][0] = 1000
print(a[0][0].item(), b[0][0].item()) # 0 1000.0
А вот во втором примере b — тоже указатель, но он направлен на совершенно другой тензор, в котором содержатся совсем иной фрагмент данных. Дело в том, что PyTorch приходится совершенно по-разному представлять в памяти float16 и int64. Следовательно, PyTorch приходится сначала скопировать данные, а затем представить их совершенно по-новому.
Теперь попробуем «на высоком уровне» объяснить, что именно представляют собой эти базовые данные. Согласно устройству PyTorch, в тензоре содержится структура данных, называемая torch.Storage, и в ней «хранятся все данные, которые просматривает torch.Tensor. Также известно, что в тензоре обязательно должна храниться и другая информация, в частности, градиент, функция обратного распространения и т.д. Можно сказать, что теперь у нас готов скелет, очертания которого напоминают torch.Tensor.
Так начинаем в первом приближении понимать, что за структура данных лежит в основе torch.Tensor.
struct Tensor {
data: torch.Storage.
grad: Tensor,
grad_fn: BackwardFn
.... // другие метаданные
}Реализуем хранение данных
Теперь давайте разберёмся, как реализуется torch.Storage. Для этого сначала рассмотрим в качестве примера простую двумерную матрицу. Тривиальнее всего было бы определить вложенный массив m = [[1,2],[3,4]]. Притом, что в таком представлении данные отражаются достаточно точно и наглядно, в данном случае возникают две основные проблемы:
- Довольно сложно работать с массивами, каждый из которых может иметь любую (произвольную) глубину из-за собственной вложенности — в особенности, на статически типизированных языках.
- Ещё важнее, что при этом возникают серьёзные накладные расходы, влияющие на производительность.
Допустим, мы хотим обратиться к элементу 2, расположенному в позиции (1,2). Для этого нам понадобится проиндексировать массив при помощи m[0][1]. Обратите внимание: m — это список списков, а значит, для обращения к любому элементу требуется выполнить две операции доступа. В более общем случае вложенная схема для работы с любым n-мерным массивом требует n обращений к памяти применительно к каждому отдельному элементу, а это не идеально.
Можно ли это оптимизировать и обойтись всего 1 обращением к памяти на каждый элемент? В данном случае наиболее важно, что схема координат (x, y) полезна нам лишь тем, что помогает визуализировать компоновку матрицы, а компьютер может обойтись и без этого. На самом деле, в соответствии с документацией по torch.Storage, сам torch.Storage – это «непрерывный одномерный массив элементов». Из этого следует, что, если мы сможем определить некую функцию, которая будет отображать на это одномерное представление нашу схему координат (x, y) (а в сущности — и любую n-мерную схему), то как раз добьёмся, чтобы для доступа к любому элементу требовалось всего один раз обратиться к памяти.
У такого представления есть несколько достоинств. С практической точки зрения мы не только получаем возможность обходиться единственным обращением к памяти на каждый элемент, но и гарантированно добиваемся смежности элементов в памяти. Поэтому мы можем эффективнее использовать пространство, доступное нам в кэше.
С точки зрения реализации становится исключительно легко понять, как именно работают операции вида .reshape() и .view(). Базовое хранилище данных не меняется, нам просто нужно адаптировать нашу функцию отображения, переведя её из координатной системы старой формы в систему новой формы. Если функция отображения уже принимает форму системы координат в качестве входного параметра, то остаётся всего лишь изменить атрибут формы.
В r-nn я создал функцию, явно отображающую индекс на одномерное представление всякий раз, когда происходит операция индексирования — сделал так для ясности. Вы можете подумать, что, вычисляя таким образом отображение всякий раз при обращении к элементу, мы несём ненужные издержки — и будете правы! В PyTorch такое отображение поддерживается при помощи атрибута stride, обновляемого всякий раз при изменении формы тензора.
В вышеприведённом примере показано, как шаг (stride) обеспечивает эффективное индексирование данных. i-й элемент шага соответствует тому, на сколько пробелов вдоль массива мы должны продвинуться, чтобы получить следующий элемент с размерностью i. При вызове .transpose() PyTorch только и требуется, что обновить шаг. Как видим, хотя мы и получили в ответ тензор в изменённом представлении, сами данные, лежащие в основе этого тензора, остались прежними. Чтобы усвоить этот момент, задумайтесь, почему приведённый ниже фрагмент кода не будет работать.
x = torch.arange(9).reshape(3, 3) # 3 x 3
x.t().view(1, -1) # Уменьшаем размерность тензора
# >> RuntimeError: размер представления несовместим с размером и шагом полученного на вход тензора
# Вместо этого воспользуйтесь .reshape()
Ошибка свидетельствует о несовпадении размеров представления с размером и шагом входного тензора, поскольку тензор был транспонирован. Это логично, поскольку невозможно уменьшить размерность тензора, не меняя заложенных в нём данных. Также на этом примере демонстрируется разница между .view() и .reshape(). .view() — это представление данных, с которыми мы работаем, тогда как .reshape() может вернуть копию данных, если формы окажутся несовместимы.
Так мы подходим к интуитивному пониманию следующего момента: форма тензора не зависит от заложенных в него данных.
struct Tensor {
data: torch.Storage.
grad: Tensor,
grad_fn: BackwardFn
shape: Array // просто атрибут
stride: Array
.... // другие метаданные
}Широковещание
Одна из основных зон, где в PyTorch возникают баги, связана с неправильной организацией широковещания. Для начала давайте вспомним правила широковещания в PyTorch:
- Оба тензора должны иметь как минимум по 1 измерению
- Все измерения, начиная с самых правых, должны иметь одинаковые размеры, причём, каждое из них должно быть либо равно 1, либо отсутствовать. На первый взгляд это может показаться путаным, поэтому проще сразу рассмотреть, как рассчитывается результирующая форма. Алгоритм строится так:
for each dimension, starting from the right:
if both shapes have this dimension:
if they are different:
neither is 1: error
else: use larger dimension
else they are the same: use dimension
else:
use whichever dimension exists
Чтобы понять, почему ситуация обстоит именно так, вспомните о рассмотренной ранее функции, которая отображает координаты на одномерный индекс. Предположим, что мы сообщим этой функции как точную форму тензора, так и форму, переданную в широковещательном режиме (при этом также будем считать, что широковещание выполнено правильно). Когда станем перебирать измерения справа налево, нам может попасться actual_shape[i] != broadcasted_shape[i], и в таком случае мы будем уверены, что данное измерение передано широковещательно. Тогда мы сможем просто вернуть 0-й элемент вдоль этого i-го измерения.
Продумаем ещё один пример. У нас есть тензор a формы 1 x 2, и мы складываем его с тензором b формы 2 x 2. Из правил широковещания нам известно, что такое возможно, и что результирующий тензор будет иметь форму 2 x 2. Можно пошагово просмотреть, как вычисляется каждый элемент.
a = torch.tensor([1, 2]).reshape(1, 2) # 1 x 2
b = torch.tensor([[3, 4], [5, 6]]) # 2 x 2
c = torch.zeros((2, 2)) # известно, что a + b даст нам тензор 2x2
# Теперь обрабатываем: первое измерение справа
c[0][0] = a[0][0] + b[0][0]
c[0][1] = a[0][1] + b[0][1]
# Теперь обрабатываем: второе измерение справа
# обратите внимание: оно было широковещательно передано в 2x2 вдоль второго измерения
# убеждаемся, что broadcasted_shape[1] = 2 что не равно
# исходному shape[1] = 1, поэтому мы знаем, что это измерение было передано широковещательно
# Следовательно, сейчас мы искусственно вернули 0-й элемент вдоль этого измерения
c[1][0] = a[0][0] + b[1][0]
c[1][1] = a[0][1] + b[1][1]
assert torch.equal(a + b, c) # true Ладно, этот пример может показаться слегка переусложнённым, и вы уже могли задуматься, почему же нужно понимать широковещание с такой степенью детализации. Дело в том, что, опираясь на знания о широковещании, мы разберёмся с двумя следующими концепциями, которые в ином случае сразу не очевидны.
Во-первых, здесь не клонируются данные. То есть, не возникает никаких издержек ни с памятью, ни с производительностью, если маленький тензор широковещательно передаётся в гораздо более крупную форму.
Во-вторых, поскольку в меньшем тензоре мы используем тот же элемент, градиенты накапливаются по элементам в этом меньшем измерении. Это особенно полезно при отладке градиентов или реализации пользовательских функций автоматического дифференцирования, при работе с которыми применяется широковещание.
Широковещание при перемножении матриц
Вполне понятно, как перемножить 2 матрицы. Но что происходит при умножении любой общей n-мерной матрицы? В PyTorch этот вопрос отлично документирован, и я всем настоятельно рекомендую почитать, что написано по ссылке выше. В частности, ситуацию для более сложного n-мерного случая я резюмировал в виде приведённого ниже алгоритма:
- Берём 2 последних измерения обоих тензоров и проверяем, можно ли их перемножить. Если нет — выдаётся ошибка.
- Широковещательно передаём оставшиеся измерения. Затем поступаем с результирующей формой так: [передано широковещательно] + [форма, полученная в результате перемножения матриц]
- 3Считаем [передано широковещательно] пакетным измерением и выполняем пакетное умножение матриц and perform batched matrix multiplication.
Снова давайте рассмотрим пример.
a = torch.randn((3, 4, 1, 2)) # 3 x 4 x 1 x 2
b = torch.randn((1, 2, 3)) # 1 x 2 x 3
# Форма при перемножении матриц: 1x2 @ 2x3 -> 1x3
# Пакетная форма: мы широковещательно передаём (3, 4) и (1) -> (3, 4)
# Результирующая форма: 3 x 4 x 1 x 3
c = torch.zeros((3, 4, 1, 3))
# перебираем пакетные измерения (3, 4)
for i in range(3):
for j in range(4):
a_slice = a[i][j] # 1 x 2
b_slice = b[0] # 2 x 3
c[i][j] = a_slice @ b_slice # 1 x 3
assert torch.equal(torch.matmul(a, b), c)Часто бывает сложно визуализировать и понять, как именно взаимодействуют друг с другом тензоры высоких размерностей. Но я думаю, что, если понять перемножение матриц по такому пакетному принципу, то станет очевиднее, чем именно нам полезен PyTorch. Само обучение будет для вас не таким утомительным, если пока вы новичок в работе с PyTorch.
Обратное распространение
Перед чтением этого раздела настоятельно рекомендую вам составить представление о том, как действует автоматическое дифференцирование выражений. Если вы этого пока не знаете, что рекомендую сначала посмотреть руководство на эту тему от Андрея Карпатого.
Суть PyTorch заложена в его движке для автоматического дифференцирования. В целом, всякий раз, когда между двумя тензорами происходит дифференцируемая операция, PyTorch автоматически выстроит через функцию обратного вызова весь её вычислительный граф. Затем градиент каждого тензора будет обновляться всякий раз после вызова функции .backward(). Это самая крупная абстракция в PyTorch. Часто бывает так, что мы вызываем .backward() и просто надеемся, что градиенты потекут правильно. В этом разделе я постараюсь помочь вам понять этот процесс и научу визуализировать потоки градиентов.
Эти 94 строки кода [в micrograd] — всё, что нам требуется для обучения нейронной сети. Всё прочее связано лишь с повышением эффективности.
Андрей Карпатый
При попытке понять обратное распространение часто возникает такая проблема: большинство людей умеют обращаться с производными и цепным правилом дифференцирования сложной функции при работе со скалярами. Но не так очевидно, как перенести это правило на работу с тензорами высоких размерностей.
Оказывается, что при создании r-nn будет достаточным просто определить производные для скалярных значений — и они обобщаются до тензоров с произвольно высокой размерностью. Поэтому я решил рассказать об обратном распространении с точки зрения работы со скалярными значениями, это очень полезно и в образовательном смысле. Далее до конца этого раздела будем считать, что каждый тензор состоит из отдельных тенхоров, каждый из которых имеет скалярное значение, собственный градиент и свою функцию обратного распространения.
Сначала познакомимся с производными сложения +, умножения ∗, возведения в степень xk, потенцирования ex и натурального логарифма ln(). 5 этих простейших действий достаточно для выполнения подавляющего большинства операций. Можете сами в этом убедиться на примере операции softmax, которая на первый взгляд может показаться сложной.
Производная от ex — это просто ex. В знаменателе происходит сложение, а деление можно представить как сочетание умножения и возведения знаменателя в степень x(-1).
Создавая нашу r-nn, мы указываем шаги обратного распространения только для этих скалярных операций, из чего естественным образом следует производная от умножения двух матриц друг на друга. Это не должно удивлять, учитывая, что мы понимаем перемножение матриц как последовательность операций умножения и сложения над множеством скалярных значений. Но я вижу определённую красоту в этом простом факте, и именно он помогает намного лучше и легче усвоить суть обратного распространения.
Размышляя о градиентах с такой локальной «скалярной» точки зрения, мы дополнительно выигрываем, так как нам становится проще усвоить, как тензорные операции выполняются над градиентами. Например, такие операции как .reshape(), .transpose(), .cat и .split() в локальном масштабе не влияют ни на отдельное значение, ни на его градиент. Из этого естественным образом следует, что эффектом воздействия таких операций на градиент тензора являются сами эти операции. Например, когда мы уменьшаем размерность тензора при помощи .reshape(-1), мы получаем точно такой же эффект, как если бы вызвали .reshape(-1) применительно к градиенту тензора.
Оптимизации
Проект r-nn написан в образовательных целях, и поэтому, готовя его, я стремился к ясности, а не к оптимизации. Несмотря на это, думаю, что было бы интересно включить в эту статью ещё один раздел, в котором были бы описаны возможные оптимизации.
Перемножение матриц
Перемножение матриц повсюду применяется в нейронных сетях (и в технологии глубокого обучения вообще), оптимизация такого перемножения — это бизнес на триллионы долларов. Я затрону здесь некоторые возможные оптимизации, осуществимые даже без GPU.
В одном из вариантов оптимизации мы опираемся на паттерны доступа к памяти, а сам алгоритм не меняем. Напомню: имея A @ B, мы раз за разом вычисляем скалярное произведение ряда в A и столбца в B. Также напомню, что в PyTorch данные хранятся вдоль столбцов. Притом, что в такой ситуации очень удобно обращаться к строкам в A, мы не добиваемся пространственной компактности кэша в B. В данном случае есть простое решение — транспонировать матрицу так, чтобы она упорядочивалась вдоль столбцов. В таком случае, загружая данные из памяти, мы сможем сразу загружать в столбец B именно те элементы, которые нужно, не меняя при этом кэш-линию. Транспонирование — это операция со сложностью O(N), поэтому её целесообразно применять только с большими матрицами. Реализация такого подхода с r-nn приводится здесь.
Ещё один секретный алгоритм заключается в следующем: можно оперировать не всей матрицей сразу, а лишь отдельными её блоками. Такой метод называется блочное матричное умножение. Идея в том, чтобы разбить матрицу на сравнительно мелкие блоки, и уже над ними выполнять операции умножения. В данном случае мы дополнительно выигрываем потому, что снижаем число промахов кэша, тоже потому, что перешли к работе с мелкими участками матрицы.
Ещё один секретный алгоритм заключается в следующем: можно оперировать не всей матрицей сразу, а лишь отдельными её блоками. Такой метод называется блочное матричное умножение. Идея в том, чтобы разбить матрицу на сравнительно мелкие блоки, и уже над ними выполнять операции умножения. В данном случае мы дополнительно выигрываем потому, что снижаем число промахов кэша, тоже потому, что перешли к работе с мелкими участками матрицы.
Память и промежуточные активации
Рассмотрим выражение (a+b)∗(c+d)=e. Под капотом Python интерпретирует это выражение так:
_t1 = a + b
_t2 = c + d
e = _t1 + _t2
При вызове e.backward() градиент должен проследовать через _t1 и _t2 перед тем, как дойдёт до листовых значений a, b, c и d. В упрощённой реализации autograd сохранятся _t1 и _t2, а при сравнительно больших вычислениях это влечёт серьёзные издержки. Однако, обратите внимание: 
, и это означает, что _t1 и _t2 нам далее не нужны. Поэтому PyTorch далее не будет хранить эти значения.
Почему это важно
Надеюсь, что смог вас убедить: когда реализуешь подобное решение с нуля, это невероятно полезно для самообразования, причём, именно таков наилучший путь к интуитивному усвоению темы, начиная от самых азов. В любом случае, надеюсь, что после изучения этого поста вы стали увереннее понимать, как именно PyTorch работает под капотом. В следующий раз, когда будете выполнять какую-нибудь операцию или вызывать .backward(), она уже не будет для вас секретом.
Весь код для r-nn выложен здесь.
Автор: ph_piter
