TensorFlow.js에서 사용자 정의 작업, 커널 및 그래디언트 작성

개요

이 가이드에서는 TensorFlow.js에서 사용자 정의 작업(ops), 커널 및 그라데이션을 정의하는 메커니즘을 간략하게 설명합니다. 주요 개념에 대한 개요와 실제 개념을 보여주는 코드에 대한 포인터를 제공하는 것을 목표로 합니다.

이 가이드는 누구를 위한 것인가요?

이것은 TensorFlow.js의 일부 내부를 다루는 상당히 고급 가이드이며, 다음 그룹의 사람들에게 특히 유용할 수 있습니다.

  • 다양한 수학 연산의 동작을 사용자 정의하는 데 관심이 있는 TensorFlow.js의 고급 사용자(예: 기존 그래디언트 구현을 재정의하는 연구자 또는 라이브러리에서 누락된 기능을 패치해야 하는 사용자)
  • TensorFlow.js를 확장하는 라이브러리를 구축하는 사용자(예: TensorFlow.js 기본 요소 또는 새로운 TensorFlow.js 백엔드 위에 구축된 일반 선형 대수 라이브러리).
  • 이러한 메커니즘이 어떻게 작동하는지에 대한 일반적인 개요를 얻고 싶은 tensorflow.js에 새로운 작업을 기여하는 데 관심이 있는 사용자입니다.

이는 내부 구현 메커니즘에 적용되므로 TensorFlow.js의 일반적인 사용에 대한 가이드가 아닙니다 . TensorFlow.js를 사용하기 위해 이러한 메커니즘을 이해할 필요는 없습니다.

이 가이드를 최대한 활용하려면 TensorFlow.js 소스 코드를 읽는 데 능숙해야 합니다(또는 기꺼이 읽어야 합니다).

술어

이 가이드에서는 몇 가지 주요 용어를 미리 설명하는 데 유용합니다.

연산(Ops) — 하나 이상의 텐서를 출력으로 생성하는 하나 이상의 텐서에 대한 수학적 연산입니다. Ops는 '상위 수준' 코드이며 다른 Ops를 사용하여 해당 논리를 정의할 수 있습니다.

커널 — 특정 하드웨어/플랫폼 기능에 연결된 특정 작업 구현입니다. 커널은 '낮은 수준'이며 백엔드에 따라 다릅니다. 일부 작업은 작업에서 커널로 일대일 매핑을 수행하는 반면 다른 작업은 여러 커널을 사용합니다.

Gradient / GradFunc — 일부 입력과 관련하여 해당 함수의 미분을 계산하는 연산/커널 의 '역방향 모드' 정의입니다. 그라데이션은 '상위 수준' 코드(백엔드에 국한되지 않음)이며 다른 작업이나 커널을 호출할 수 있습니다.

커널 레지스트리 - (커널 이름, 백엔드 이름) 튜플에서 커널 구현으로의 맵입니다.

Gradient Registry커널 이름에서 그래디언트 구현까지의 맵입니다.

코드 구성

작업그라데이션은 tfjs-core 에 정의되어 있습니다.

커널은 백엔드별로 다르며 해당 백엔드 폴더(예: tfjs-backend-cpu )에 정의됩니다.

이러한 패키지 내에서는 사용자 정의 작업, 커널 및 그라데이션을 정의할 필요가 없습니다. 그러나 구현 시 유사한 기호를 자주 사용합니다.

맞춤 작업 구현

사용자 정의 작업을 생각하는 한 가지 방법은 종종 텐서를 입력으로 사용하여 일부 텐서 출력을 반환하는 JavaScript 함수와 같습니다.

  • 일부 작업은 기존 작업의 관점에서 완전히 정의될 수 있으며 이러한 함수를 직접 가져오고 호출해야 합니다. 여기에 예가 있습니다 .
  • 작업 구현은 백엔드 특정 커널로 디스패치될 수도 있습니다. 이는 Engine.runKernel 통해 수행되며 "사용자 정의 커널 구현" 섹션에서 자세히 설명합니다. 여기에 예가 있습니다 .

사용자 정의 커널 구현

백엔드별 커널 구현을 통해 특정 작업에 대한 논리의 최적화된 구현이 가능합니다. 커널은 tf.engine().runKernel() 호출하는 작업에 의해 호출됩니다. 커널 구현은 네 가지로 정의됩니다.

  • 커널 이름.
  • 커널이 구현되는 백엔드입니다.
  • 입력: 커널 함수에 대한 텐서 인수.
  • 속성: 커널 함수에 대한 비텐서 인수.

다음은 커널 구현 의 예입니다. 구현에 사용되는 규칙은 백엔드마다 다르며 각 특정 백엔드의 구현 및 문서를 살펴보면 가장 잘 이해됩니다.

일반적으로 커널은 텐서보다 낮은 수준에서 작동하며 대신 tfjs-core에 의해 결국 텐서로 래핑될 메모리에 직접 읽고 씁니다.

커널이 구현되면 tfjs-core의 registerKernel 함수를 사용하여 TensorFlow.js에 등록할 수 있습니다. 커널이 작동하기를 원하는 모든 백엔드에 대해 커널을 등록할 수 있습니다. 일단 등록되면 커널은 tf.engine().runKernel(...) 을 사용하여 호출할 수 있으며 TensorFlow.js는 다음의 구현으로 디스패치합니다. 현재 활성 백엔드.

사용자 정의 그라데이션 구현

그라데이션은 일반적으로 주어진 커널( tf.engine().runKernel(...) 호출에 사용된 동일한 커널 이름으로 식별됨)에 대해 정의됩니다. 이를 통해 tfjs-core는 레지스트리를 사용하여 런타임 시 모든 커널에 대한 그래디언트 정의를 조회할 수 있습니다.

사용자 정의 그라데이션 구현은 다음과 같은 경우에 유용합니다.

  • 라이브러리에 없을 수 있는 그라데이션 정의 추가
  • 주어진 커널에 대한 기울기 계산을 사용자 정의하기 위해 기존 기울기 정의를 재정의합니다.

여기에서 그라데이션 구현 의 예를 볼 수 있습니다.

특정 호출에 대한 그래디언트를 구현한 후에는 tfjs-core의 registerGradient 함수를 사용하여 TensorFlow.js에 등록할 수 있습니다.

그래디언트 레지스트리를 우회하여 임의의 방식으로 임의 함수에 대한 그래디언트를 계산할 수 있는 사용자 정의 그래디언트를 구현하는 또 다른 접근 방식은 tf.customGrad 를 사용하는 것입니다.

다음은 customGrad를 사용하는 라이브러리 내 작업의 예 입니다.