Шта је Гоогле ЈАКС? Све што треба да знате

Гоогле ЈАКС или Јуст Афтер Екецутион је оквир који је развио Гоогле да убрза задатке машинског учења.

Можете га сматрати библиотеком за Питхон, која помаже у бржем извршавању задатака, научном рачунарству, трансформацијама функција, дубоком учењу, неуронским мрежама и још много тога.

О Гоогле ЈАКС-у

Најосновнији рачунски пакет у Питхон-у је пакет НумПи који има све функције као што су агрегације, векторске операције, линеарна алгебра, н-димензионалне манипулације низом и матрицом и многе друге напредне функције.

Шта ако бисмо могли додатно да убрзамо прорачуне изведене помоћу НумПи-а – посебно за огромне скупове података?

Да ли имамо нешто што би могло да ради подједнако добро на различитим типовима процесора као што су ГПУ или ТПУ, без икаквих промена кода?

Шта кажете на то да систем може аутоматски и ефикасније да изврши трансформације функција које се могу саставити?

Гоогле ЈАКС је библиотека (или оквир, како каже Википедија) која ради управо то, а можда и много више. Направљен је да оптимизује перформансе и ефикасно обавља задатке машинског учења (МЛ) и дубоког учења. Гоогле ЈАКС пружа следеће функције трансформације које га чине јединственим од других МЛ библиотека и помажу у напредном научном рачунању за дубоко учење и неуронске мреже:

• Ауто диференцијација
• Ауто векторизација
• Аутоматска паралелизација
• Јуст-ин-тиме (ЈИТ) компилација

Јединствене функције Гоогле ЈАКС-а

Све трансформације користе КСЛА (убрзану линеарну алгебру) за веће перформансе и оптимизацију меморије. КСЛА је машина за оптимизацију компајлера специфична за домен који изводи линеарну алгебру и убрзава ТенсорФлов моделе. Коришћење КСЛА на врху вашег Питхон кода не захтева значајне промене кода!

Хајде да детаљно истражимо сваку од ових карактеристика.

Карактеристике Гоогле ЈАКС-а

Гоогле ЈАКС долази са важним функцијама трансформације за састављање ради побољшања перформанси и ефикаснијег обављања задатака дубоког учења. На пример, аутоматско диференцирање да бисте добили градијент функције и пронашли деривате било ког реда. Слично, аутоматска паралелизација и ЈИТ за паралелно обављање више задатака. Ове трансформације су кључне за апликације попут роботике, игара, па чак и истраживања.

Функција састављања трансформације је чиста функција која трансформише скуп података у други облик. Називају се састављајућим јер су самосталне (тј. ове функције немају зависности од остатка програма) и немају стање (тј. исти улаз ће увек резултирати истим излазом).

И(к) = Т: (ф(к))

У горњој једначини, ф(к) је оригинална функција на коју се примењује трансформација. И(к) је резултујућа функција након примене трансформације.

На пример, ако имате функцију под називом ‘тотал_билл_амт’, и желите резултат као трансформацију функције, можете једноставно користити трансформацију коју желите, рецимо градијент (град):

град_тотал_билл = град(укупни_рачун_амт)

Трансформисањем нумеричких функција коришћењем функција као што је град(), лако можемо добити њихове деривате вишег реда, које можемо интензивно користити у алгоритмима за оптимизацију дубоког учења као што је спуштање градијента, чинећи тако алгоритме бржим и ефикаснијим. Слично, коришћењем јит(), можемо компајлирати Питхон програме тачно на време (лењо).

#1. Ауто диференцијација

Питхон користи функцију аутоград да аутоматски разликује НумПи и изворни Питхон код. ЈАКС користи модификовану верзију аутоград (тј. град) и комбинује КСЛА (убрзану линеарну алгебру) да изврши аутоматску диференцијацију и пронађе деривате било ког редоследа за ГПУ (Графичке процесорске јединице) и ТПУ (Тензорске процесорске јединице).]

Кратка напомена о ТПУ-у, ГПУ-у и ЦПУ-у: ЦПУ или централна процесорска јединица управља свим операцијама на рачунару. ГПУ је додатни процесор који побољшава рачунарску снагу и покреће врхунске операције. ТПУ је моћна јединица посебно развијена за сложена и тешка радна оптерећења попут вештачке интелигенције и алгоритама дубоког учења.

На исти начин као и функција аутоград, која се може разликовати кроз петље, рекурзије, гране и тако даље, ЈАКС користи функцију град() за градијенте обрнутог режима (повратно ширење). Такође, можемо разликовати функцију од било ког реда користећи град:

град(град(град(син θ))) (1.0)

Ауто диференцијација вишег реда

Као што смо раније споменули, град је прилично користан у проналажењу парцијалних извода функције. Можемо користити делимични извод за израчунавање градијента спуштања функције трошкова у односу на параметре неуронске мреже у дубоком учењу да бисмо минимизирали губитке.

Израчунавање парцијалног извода

Претпоставимо да функција има више променљивих, к, и и з. Проналажење извода једне променљиве задржавањем константних осталих променљивих назива се делимични извод. Претпоставимо да имамо функцију,

ф(к,и,з) = к + 2и + з2

Пример за приказ делимичног извода

Делимични извод к ће бити ∂ф/∂к, што нам говори како се функција мења за променљиву када су друге константне. Ако ово изведемо ручно, морамо написати програм за разликовање, применити га за сваку променљиву, а затим израчунати спуштање градијента. Ово би постала сложена и дуготрајна ствар за више варијабли.

Аутоматско диференцијација разлаже функцију на скуп елементарних операција, као што су +, -, *, / или син, цос, тан, екп, итд., а затим примењује ланчано правило за израчунавање извода. То можемо да урадимо иу напред и назад режиму.

Ово није то! Сви ови прорачуни се дешавају тако брзо (па, размислите о милион калкулација сличних горе наведеном и времену које може потрајати!). КСЛА брине о брзини и перформансама.

#2. Убрзана линеарна алгебра

Узмимо претходну једначину. Без КСЛА, израчунавање ће трајати три (или више) језгра, где ће свако језгро обављати мањи задатак. На пример,

Кернел к1 –> к * 2и (множење)

к2 –> к * 2и + з (сабирање)

к3 –> Смањење

Ако исти задатак обавља КСЛА, једно језгро се брине за све међуоперације тако што их спаја. Средњи резултати елементарних операција се стримују уместо да се чувају у меморији, чиме се штеди меморија и повећава брзина.

#3. Компилација тачно на време

ЈАКС интерно користи КСЛА компајлер за повећање брзине извршавања. КСЛА може повећати брзину ЦПУ-а, ГПУ-а и ТПУ-а. Све ово је могуће коришћењем извршавања ЈИТ кода. Да бисмо ово користили, можемо користити јит путем увоза:

from jax import jit
def my_function(x):
	…………some lines of code
my_function_jit = jit(my_function)

Други начин је украшавање јит-а преко дефиниције функције:

@jit
def my_function(x):
	…………some lines of code

Овај код је много бржи јер ће трансформација вратити преведену верзију кода позиваоцу уместо да користи Питхон интерпретер. Ово је посебно корисно за векторске улазе, као што су низови и матрице.

Исто важи и за све постојеће Питхон функције. На пример, функције из пакета НумПи. У овом случају, требало би да увеземо јак.нумпи као јнп уместо НумПи:

import jax
import jax.numpy as jnp

x = jnp.array([[1,2,3,4], [5,6,7,8]])

Када то урадите, основни објекат ЈАКС низа који се зове ДевицеАрраи замењује стандардни НумПи низ. ДевицеАрраи је лењ – вредности се чувају у акцелератору док не буду потребне. Ово такође значи да ЈАКС програм не чека да се резултати врате у позивајући (Питхон) програм, пратећи тако асинхроно слање.

#4. аутоматска векторизација (вмап)

У типичном свету машинског учења, имамо скупове података са милион или више тачака података. Највероватније бисмо извршили неке прорачуне или манипулације на свакој или већини ових тачака података – што је задатак који одузима много времена и меморије! На пример, ако желите да пронађете квадрат сваке тачке података у скупу података, прва ствар на коју бисте помислили је да направите петљу и узмете квадрат један по један – аргх!

Ако креирамо ове тачке као векторе, могли бисмо да урадимо све квадрате одједном вршећи векторске или матричне манипулације на тачкама података помоћу нашег омиљеног НумПи-а. А ако би ваш програм могао то да уради аутоматски – можете ли тражити нешто више? То је управо оно што ЈАКС ради! Може аутоматски да векторизује све ваше тачке података тако да можете лако да извршите све операције на њима – чинећи ваше алгоритме много бржим и ефикаснијим.

ЈАКС користи функцију вмап за аутоматску векторизацију. Размотрите следећи низ:

x = jnp.array([1,2,3,4,5,6,7,8,9,10])
y = jnp.square(x)

Радећи само горе наведено, квадрат метода ће се извршити за сваку тачку у низу. Али ако урадите следеће:

vmap(jnp.square(x))

Квадрат методе ће се извршити само једном јер се тачке података сада аутоматски векторизују помоћу вмап методе пре извршавања функције, а петља се гура надоле на елементарни ниво операције – што резултира множењем матрице уместо скаларног множења, што даје боље перформансе .

#5. СПМД програмирање (пмап)

СПМД – или једнопрограмско програмирање више података је неопходно у контексту дубоког учења – често бисте применили исте функције на различите скупове података који се налазе на више ГПУ-а или ТПУ-ова. ЈАКС има функцију под називом пумпа, која омогућава паралелно програмирање на више ГПУ-а или било ком акцелератору. Као и ЈИТ, КСЛА ће компајлирати програме који користе пмап и извршавати их истовремено у свим системима. Ова аутоматска паралелизација ради и за унапред и за реверзна израчунавања.

Како функционише пмап

Такође можемо применити више трансформација одједном у било ком редоследу на било коју функцију као:

пмап(вмап(јит(град (ф(к)))))

Вишеструке композитне трансформације

Ограничења Гоогле ЈАКС-а

Гоогле ЈАКС програмери су добро размислили о убрзавању алгоритама дубоког учења док су уводили све ове сјајне трансформације. Функције и пакети научног рачунања су на линији НумПи, тако да не морате да бринете о криву учења. Међутим, ЈАКС има следећа ограничења:

• Гоогле ЈАКС је још увек у раним фазама развоја, и иако је његова главна сврха оптимизација перформанси, не пружа много користи за ЦПУ рачунарство. Чини се да НумПи ради боље, а коришћење ЈАКС-а може само повећати трошкове.
• ЈАКС је још увек у свом истраживању или раној фази и потребно му је више финог подешавања да би достигао инфраструктурне стандарде оквира као што је ТенсорФлов, који су више успостављени и имају више унапред дефинисаних модела, пројеката отвореног кода и материјала за учење.
• За сада, ЈАКС не подржава Виндовс оперативни систем – потребна вам је виртуелна машина да би функционисала.
• ЈАКС ради само на чистим функцијама – онима које немају никакве нуспојаве. За функције са нежељеним ефектима, ЈАКС можда није добра опција.

Како да инсталирате ЈАКС у вашем Питхон окружењу

Ако имате Питхон подешавање на вашем систему и желите да покренете ЈАКС на вашој локалној машини (ЦПУ), користите следеће команде:

pip install --upgrade pip
pip install --upgrade "jax[cpu]"

Ако желите да покренете Гоогле ЈАКС на ГПУ или ТПУ, пратите упутства дата на ГитХуб ЈАКС страна. Да бисте подесили Питхон, посетите питхон званична преузимања страна.

Закључак

Гоогле ЈАКС је одличан за писање ефикасних алгоритама дубоког учења, роботике и истраживања. Упркос ограничењима, он се интензивно користи са другим оквирима као што су Хаику, Флак и многи други. Моћи ћете да цените шта ЈАКС ради када покренете програме и видите временске разлике у извршавању кода са и без ЈАКС-а. Можете почети читањем званична Гоогле ЈАКС документацијашто је прилично свеобухватно.