ML: Вступ до PyTorch: 2. Графи

Вступ
Прямий і зворотний проходи
Побудова графа
Обчислення градієнтів
Градієнт у проміжних вузлах
Призупинення побудови графа

Від’єднання вузла від графа
Трохи прикладів
Що не можна робити з листками графа
Зрізи тензорів
Візуалізація
Мінімізація функції

Вступ

У цьому документі міститься продовження опису роботи з тензорами у PyTorch. Тепер ми розглянемо обчислювальні графи і поширення на них градієнтів. Наступний документ присвячений основам роботи з нейронними мережами.

Прямий і зворотний проходи

Нагадаємо, що обчислювальний граф є послідовністю дій для отримання значення деякої величини. У машинному навчанні зазвичай це скалярна функція помилки (тензор нульової розмірності). Результат обчислення отримується при прямому проході по графу. Нижче намальований граф для функції $z = x\cdot x + \sin(2\cdot y)$ і наведені значення вузлів (зелений колір) при x=4 і y=3.14159265:

$$ z = x\cdot x + \sin(2\cdot y) $$ $$ \begin{array}{lcl} g_x = \frac{\displaystyle\partial z}{\displaystyle\partial x} &=& 2\cdot x\Bigr|_{x=4} = 8,\\[4mm] g_y = \frac{\displaystyle\partial z}{\displaystyle\partial y} &=& 2\,\cos(2\cdot y)\Bigr|_{y=\pi} = 2. \end{array} $$

Після прямого проходу по графу (від листків до кореня $z$ дерева), запускається процедура зворотного проходу, що обчислює похідні (градієнти, синій колір) цільового виразу $z$ за іншими вузлами графа. На рисунку з вузла z виходить $g_z=1$. Сумарний $g_x=8$, що входить у вузол x, дорівнює частковій похідній $\partial z/\partial x$, а $g_y=2$, що входить у y, дорівнює $\partial z/\partial y$. З деталями обчислення градієнтів варто ознайомитися в документі "Обчислювальний граф".

У PyTorch графи динамічні. Вони будуються по ходу обчислення виразів. Щоб повторити обчислення з новими даними, необхідно повторно "пройти" по всіх виразах. Цей підхід відмінний від статичних графів, які будуються в tensorflow. Статичний граф визначається один раз, компілюється і потім може довільну кількість разів запускатися на обчислення з різними значеннями в листових вузлах, але при цьому не може бути змінений.

Побудова графа

Тензор у PyTorch, окрім даних зі значеннями елементів, може зберігати градієнти за цими елементами і багато іншого, необхідного для роботи з обчислювальним графом:

from torch import tensor, empty, ones, zeros

v  = zeros(2)                # 2D - вектор з нулів

print(v.data)                # tensor([0.,0.])  - дані тензора (те саме, що просто v)
print(v.grad)                # None             - градієнт за тензором (поки його немає)
print(v.grad_fn)             # None             - функція до нього привела (поки немає графа)
print(v.is_leaf)             # True             - є листком графа (так)
print(v.requires_grad)       # False            - за ним потрібен градієнт (поки не потрібен)

Усі операції з тензором v виконуються з його атрибутом v.data. Атрибут v.grad (якщо він є) є тензором тієї самої розмірності, що й v (у нього також є data, grad,..., і т.д.).

У PyTorch граф починає будуватися, якщо у виразі є тензор з атрибутом requires_grad у значенні True.
Цей атрибут можна задати в конструкторі (при створенні тензора) або у будь-який момент пізніше:

x = ones(2, requires_grad=True)          # вектор [1., 1.] одразу буде вузлом графа

y = empty(2).fill_(3)                    # спочатку створили вектор [3.,3.],
y.requires_grad = True                   # а пізніше оголосили його вузлом

print(y)                                 # tensor([3., 3.], requires_grad=True)

Атрибут requires_grad "заразний" і якщо у виразі є хоча б один такий тензор, виникає граф. Кожен його не листовий вузол містить останню операцію, яка до нього призводить (в атрибуті grad_fn):

z = (y*y).sum()                          # скаляр (dim=0) y[0]**2 + y[1]**2
 
print(z)                                 # tensor(18., grad_fn=<SumBackward0>)
print(y.is_leaf, z.is_leaf)              # True False
print(y.requires_grad, y.grad_fn)        # True None
print(z.requires_grad, z.grad_fn)        # True <SumBackward0>

Вузол y є листком (is_leaf), тоді як z - ні (він є кореневим = фінальним вузлом дерева).

Обчислення градієнтів

Метод backward() кореневого вузла графа запускає процедуру обчислення градієнтів у листових (is_leaf) вузлах, що мають атрибут requires_grad.
Для прикладу вище число 1, проходячи через вузол суми, дублюється стільки разів, скільки було сумувань, перетворюючись на вектор [1,1]. Потім на вузлі добутку без згортки він множиться на протилежний аргумент:

print(y.grad)                              # None
z.backward()                               # запускаємо обчислення градієнтів
print(y.grad)                              # tensor([6., 6.]) - сума 2-х вхідних grad

Повторно метод backward() викликати не можна (лише знову перебудувавши граф). Виняток становить такий його виклик: z.backward(retain_graph = True). Але в цьому випадку градієнти будуть накопичуватися (сумуватися).

Листові вузли без атрибута requires_grad=True вважаються константами і за ними градієнт не обчислюється. Нижче є дві константи (8, a) :

x = ones(1, requires_grad=True)
a = tensor(-4.)
z = 8*x + a

z.backward()

print(x, x.grad)         # tensor([1.], requires_grad=True), tensor([8.])
print(a, a.grad)         # tensor([-4.]), None
print(z)                 # tensor([4.], grad_fn=<AddBackward0>)

При аналізі обчислювального графа варто пам’ятати, що будь-яка змінна (листова або проміжна) завжди представляється одним вузлом. Якщо змінна використовується в різних обчисленнях, то з неї виходить кілька ребер, по яких потім (при зворотному поширенні) входять (сумуючись) кілька градієнтів. Таким був вузол y у прикладі з початку розділу.

Градієнт у проміжних вузлах

За замовчуванням проміжні (не листові) вузли графа не зберігають градієнти, що пройшли через них. Цю ситуацію можна змінити,
викликавши для конкретного вузла метод retain_grad:

x = tensor(2., requires_grad=True)

y  = x**2;    y.retain_grad()                    
z  = 2*y;     z.retain_grad()       

z.backward()

print(z.item(),     y.item(),     x.item())      # 8.0 4.0 2.0
print(z.grad.item(),y.grad.item(),x.grad.item()) # 1.0 2.0 8.0

У цьому прикладі коренем дерева обчислень є тензор z, а єдиним листком, що вимагає градієнта за замовчуванням - вузол x. Вузол y - це проміжний вузол.

Призупинення побудови графа

Граф повинен будуватися щоразу при повторному обчисленні градієнтів:

for i in range(1,3):
    x = empty(2).fill_(i).requires_grad_(True)
    
    z = x.dot(x)                           # граф
    z.backward()                           # отримуємо градієнти
    
    print(z.item(), x.grad)

Вище в циклі два рази створюється новий тензор x, потім обчислювальний граф для отримання суми квадратів його компонент: $z=x^2_0+x^2_1$. Градієнт від листової змінної дорівнює: $\partial z/\partial x_i = 2x_i$.

Іноді з листовими вузлами необхідно проробити дії, не змінюючи при цьому графа. Такі дії проводять під оточенням no_grad, яке блокує створення нових вузлів графа. Так, у прикладі нижче під тензор x пам’ять виділяється лише один раз (для великих тензорів це важливо). Потім у циклі під no_grad відбувається зміна значень у цій пам’яті і далі будується граф. Оскільки листовий тензор x не перестворюється, у нього необхідно обнулити градієнти, інакше вони почнуть сумуватися на наступних ітераціях циклу:

x = empty(2).requires_grad_(True)

for i in range(1,3):
    with torch.no_grad():                  # disabled gradient calculation
        x.fill_(i)                         # міняємо існуючий 
        
    z = x.dot(x)
    z.backward()                           # обчислювальний граф
    
    print(z.item(), x.grad.numpy())    
    
    x.grad.zero_()                         # обнуляємо градієнти

Цей і попередній приклад призведуть до однакових результатів:

2.0  [2., 2.]
8.0  [4., 4.]

Інший спосіб зміни даних тензора без зміни обчислювального графа — це безпосередня робота з його атрибутом data. Так, вище можна було б написати:

for i in range(1,3):
    x.data.fill_(i)                         # міняємо існуючий 
...

Після закінчення оточення with з методом torch.no_grad() побудова графа увімкнеться сама. Це ж можна зробити руками за допомогою методу torch.enable_grad():

x = ones(1, requires_grad=True)
with torch.no_grad():                      # вимикаємо побудову графа
    z1 = 2 * x
    with torch.enable_grad():              # вмикаємо  побудову графа 
        z2 = 2 * x
print(x.requires_grad, z1.requires_grad, z2.requires_grad)  # True False True

Приклад обчислень у циклі для пошуку оптимальних параметрів лінійної моделі можна знайти в цьому документі.

Від’єднання вузла від графа

За допомогою методу detach можна отримати "від’єднаний" від графа тензор, який буде посилатися на дані вихідного вузла, але не бути частиною графа:

x = tensor([1.,2.], requires_grad=True)   
y = x.detach()                                

print(x)                                   # tensor([1., 2.], requires_grad=True)
print(y)                                   # tensor([1., 2.])
y[0]=5 
print(x)                                   # tensor([5., 2.], requires_grad=True)

Це альтернативний до оточення no_grad() спосіб зміни листових вузлів без зміни графа:

x  = empty(2).requires_grad_(True)
xd = x.detach()
 
for i in range(1,3):
    xd.fill_(i)
         
    z = x.dot(x)                   # починаємо будувати граф
    z.backward()                   # обчислювальний граф
     
    print(z.item(), x.grad)
     
    x.grad.zero_()                 # обнуляємо градієнти

Трохи прикладів

Розглянемо приклад, у якому при прямому проході виходить одне й те саме значення, однак градієнти виявляються різними залежно від способів формування змінної y:

x  = tensor(3.).requires_grad_(True)   #   x.grad y.grad  y.requires_grad
y  = tensor(3.).requires_grad_(True)   # A:  1    6       True  
#y = x                                 # B:  7    7       True  
#y = tensor(1.)                        # C:  1    None    False 
#y = x.detach().clone()                # D:  1    None    False
#y = x.clone()                         # E:  7    None    True, grad_fn=<CloneBackward>

z = x + y*y                            # z = 12
z.backward()

Випадок A: похідні обчислюються за двома незалежними змінними $z=z(x,y)=x+y^2$. Це призводить до градієнтів: $\nabla_x\, z(x,y) = 1$ і $\nabla_y\, z(x,y) = 2y$.
Випадок B: присвоєння y=x відбувається за посиланням і змінна y є просто "іншим позначенням" x. Тому $z=z(x)=x+x^2$ і $\nabla_x\, z(x) = 1+2x$ (у $y$ буде такий самий градієнт).
Випадки С, D еквівалентні. У методі torch.tensor(1.) створюється тензор без requires_grad, тому для графа це константа. Аналогічно, x.clone() робить копію x, яка потім методом detach "від’єднується" від графа (нижче вона знову приєднається, але вже як константа).
Випадок E трохи складніший і найменш осмислений. Метод x.clone() робить копію x. Ця операція є вузлом графа. При зворотному проході в листовому вузлі зійдуться два градієнти (від додавання і від операції клонування), які, склавшись, дадуть 7.

Насправді клонування без від’єднання іноді призводить до не найочікуванішого результату, тому під графом при клонуванні зазвичай варто робити detach().clone() або викликати його під оточенням no_grad:

x = torch.ones(1, requires_grad=True)
                               # y.requires_grad:  y.grad_fn:
with torch.no_grad():          
    y = x.clone()              # False             None
y = x.detach().clone()         # False             None
y = x.clone()                  # True              <CloneBackward>

Що не можна робити з листками

У PyTorch стартові (листові) змінні, за якими обчислюють градієнт, не повинні брати участь в in-place обчисленнях і їх не можна перезаписувати. Розглянемо ці обмеження докладніше.

Нагадаємо, що in-place операції змінюють значення змінної не створюючи нової пам’яті. У PyTorch такими є всі методи з підкресленням у кінці імені: fill_(), add_(), mm_() і т.д. Нижче в останньому рядку відбувається не in-place обчислення (результат x+1 записується в нову пам’ять - див. значення id():

x = ones(1);    print(x, id(x))   # tensor([1.]) 2769629314008

x += 1;         print(x, id(x))   # tensor([2.]) 2769629314008   in-place
x.add_(1);      print(x, id(x))   # tensor([3.]) 2769629314008   in-place
x = x + 1;      print(x, id(x))   # tensor([4.]) 2769629311928   non in-place

Наступний код призведе до помилки "a leaf Variable that requires grad has been used in an in-place operation":

x = tensor(1., requires_grad=True)
x += 1                             # in-place для листка заборонений!

Та сама помилка виникне в наступному коді (y отримує посилання на x, тому це і є x):

x = ones(1.,requires_grad=True); print(x,id(x)) # tensor(1.,requires_grad=True) ...95208
y = x                          ; print(y,id(y)) # tensor(1.,requires_grad=True) ...95208
y *= 1                                          # in-place для листка

Для нелистових вузлів використовувати in-place обчислення можна:

x = tensor(1.,requires_grad=True)   # x.grad = tensor(0.5)
y = 2*x                             # tensor(2.,     grad_fn=<MulBackward0>)
y += 2                              # tensor(4.,     grad_fn=<AddBackward0>)
y.log_()                            # tensor(1.3863, grad_fn=<LogBackward> )
y.backward()                        # y = log(2*x+2); y'=1/(x+1)

Листова змінна не повинна переприсвоюватися, бо вона буде "знищена" і втратить властивість requires_grad:

x = tensor(1., requires_grad=True) #  tensor(1., requires_grad=True)
x = x + 1                          #  tensor(2., grad_fn=<AddBackward0>) немає requires_grad

Зрізи тензорів

Зрізи (slice) тензорів повертають новий тензор з частиною даних вихідного. При цьому вони використовують для зберігання даних спільну пам’ять. Тому обчислення градієнтів на графах з функціями зрізів вимагає певної акуратності:

x = tensor(1., requires_grad=True)    
s = ones(2)
s[1] =  s[0] * x                      # s[1] =  s[0].clone() * x <- так треба!!!
z = s.sum()                           # z = s[0] + s[0]*x
z.backward()

Цей код призведе до помилки: "one of the variables needed for gradient computation has been modified by an inplace operation". Щоб її усунути, необхідно зробити копію тензора-елемента: s[1] = s[0].clone() * x. При цьому повинен використовуватися саме метод clone(): "Unlike copy_(), this function is recorded in the computation graph. Gradients propagating to the cloned tensor will propagate to the original tensor." Зокрема на графі у тензора s маємо: s.grad_fn=<CopySlices>.

Присвоєння в зріз є in-place операцією, тому заборонене для листків. Наступний код призведе до помилки:

x, w = torch.randn(1), torch.randn(1, requires_grad=True)
w[0] = 1.
y = x*w
y.backward()

Нарешті, зрізи можуть суттєво сповільнювати зворотне поширення градієнта. Нижче два коди виконують однакові обчислення, однак правий код працює майже в 10 разів повільніше:

y = []       
for i in range(100):
    x = torch.randn(1,256,256, 
                    requires_grad=True)
    y.append( x )        
y = torch.cat(y, dim=0)
z = y.sum()
z.backward()

y = torch.empty(100,256,256)        
for i in range(100):
    x = torch.randn(256,256, 
                    requires_grad=True)
    y[i] = x   
    
z = y.sum()
z.backward()

Якщо закоментувати z.backward(), то час обчислення стане порівнянним.

Таким чином, при роботі з обчислювальними графами варто уникати слайс-копіювань, що призводять до grad_fn=<CopySlices>

Візуалізація

Для візуалізації обчислювальних графів можна використовувати невелику бібліотеку torchviz (див. її документацію і приклади):

import torchviz
from torch import tensor, empty, ones, zeros

w, b = ones(5, requires_grad=True), tensor(0., requires_grad=True)
x = ones(5)
z = x.dot(w) + b

torchviz.make_dot(z,  params = {'x': x, 'w': w, 'b': b} )

Звернімо увагу, що бібліотека малює лише листові вузли, для яких установлено requires_grad=True.

Мінімізація функції

Наведемо приклад мінімізації функції кількох змінних за допомогою градієнтного методу. Для цього скористаємося стандартним оптимізатором SGD:

import torch

def fun(x):    
    return x[0]**2 + (x[1]-1)**2

x = torch.tensor([1.,2.], requires_grad=True)    # початкові значення

optimizer = torch.optim.SGD([x],  lr=1, momentum=0.5) 
for it in range(20):
    optimizer.zero_grad()                        # обнуляємо градієнти  
    y = fun(x)                                   # обчислюємо значення функції
    y.backward()                                 # обчислюємо градієнти     
    optimizer.step()                             # підправляємо параметри
    
    print(y.item(), x.detach().numpy())