Menjalankan Program JAX daripada Dart Menggunakan C FFI

? Mengapa Gabungkan Dart dan JAX untuk Pembelajaran Mesin?

Apabila membina aplikasi, pemilihan alatan yang betul adalah penting. Anda mahukan prestasi tinggi, pembangunan mudah dan penggunaan merentas platform yang lancar. Rangka kerja popular menawarkan pertukaran:

• C memberikan kelajuan tetapi boleh melambatkan pembangunan.
• Dart (dengan Flutter) lebih perlahan tetapi memudahkan pengurusan memori dan pembangunan merentas platform.

Tetapi inilah tangkapannya: kebanyakan rangka kerja tidak mempunyai sokongan pembelajaran mesin asli (ML) yang teguh. Jurang ini wujud kerana rangka kerja ini mendahului ledakan AI. Soalannya ialah:

Bagaimanakah kita boleh menyepadukan ML dengan cekap ke dalam aplikasi?

Penyelesaian biasa seperti ONNX Runtime membenarkan pengeksportan model ML untuk penyepaduan aplikasi, tetapi ia tidak dioptimumkan untuk CPU atau cukup fleksibel untuk algoritma umum.

Masukkan JAX, perpustakaan Python yang:

• Mendayakan penulisan ML yang dioptimumkan dan algoritma tujuan umum.
• Menawarkan pelaksanaan platform-agnostik pada CPU, GPU dan TPU.
• Menyokong ciri canggih seperti autograd dan Kompilasi JIT.

Dalam artikel ini, kami akan menunjukkan kepada anda cara untuk:

1. Tulis program JAX dalam Python.
2. Jana spesifikasi XLA.
3. Kerahkan kod JAX yang dioptimumkan dalam Dart menggunakan C FFI.

? Apa itu JAX?

JAX adalah seperti NumPy pada steroid. Dibangunkan oleh Google, ia merupakan pustaka peringkat rendah berprestasi tinggi yang menjadikan ML boleh diakses namun berkuasa.

• Agnostik Platform: Kod dijalankan pada CPU, GPU dan TPU tanpa pengubahsuaian.
• Kelajuan: Dikuasakan oleh pengkompil XLA, JAX mengoptimumkan dan mempercepatkan pelaksanaan.
• Fleksibiliti: Sesuai untuk model ML dan algoritma umum.

Berikut ialah contoh membandingkan NumPy dan JAX:

# NumPy version
import numpy as np  
def assign_numpy():  
  a = np.empty(1000000)  
  a[:] = 1  
  return a  

# JAX version
import jax.numpy as jnp  
import jax  

@jax.jit  
def assign_jax():  
  a = jnp.empty(1000000)  
  return a.at[:].set(1)  

Penanda aras dalam Google Colab mendedahkan kelebihan prestasi JAX:

• CPU & GPU: JAX lebih pantas daripada NumPy.
• TPU: Peningkatan kelajuan menjadi ketara untuk model besar disebabkan kos pemindahan data.

Fleksibiliti dan kelajuan ini menjadikan JAX sesuai untuk persekitaran pengeluaran yang prestasi adalah penting.

?️ Membawa JAX ke dalam Pengeluaran

Cloud Microservices vs. Local Deployment

• Cloud: Perkhidmatan mikro Python dalam kontena bagus untuk pengiraan berasaskan awan.
• Tempatan: Menghantar penterjemah Python tidak sesuai untuk apl tempatan.

Penyelesaian: Manfaatkan Kompilasi XLA JAX

JAX menterjemah kod Python ke dalam spesifikasi HLO (Pengoptimum Tahap Tinggi), yang boleh disusun dan dilaksanakan menggunakan pustaka C XLA. Ini membolehkan:

1. Menulis algoritma dalam Python.
2. Menjalankannya secara asli melalui perpustakaan C.
3. Menyepadukan dengan Dart melalui FFI (Antara Muka Fungsi Asing).

✍️ Penyepaduan Langkah demi Langkah

1. Hasilkan Proto HLO

Tulis fungsi JAX anda dan eksport perwakilan HLOnya. Contohnya:

# NumPy version
import numpy as np  
def assign_numpy():  
  a = np.empty(1000000)  
  a[:] = 1  
  return a  

# JAX version
import jax.numpy as jnp  
import jax  

@jax.jit  
def assign_jax():  
  a = jnp.empty(1000000)  
  return a.at[:].set(1)  

Untuk menjana HLO, gunakan skrip jax_to_ir.py daripada repositori JAX:

import jax.numpy as jnp  

def fn(x, y, z):  
  return jnp.dot(x, y) / z  

Letakkan fail yang terhasil (fn_hlo.txt dan fn_hlo.pb) dalam direktori aset apl anda.

2. Bina Perpustakaan Dinamik C

Ubah suai Kod Contoh C JAX

Klon repositori JAX dan navigasi ke jax/examples/jax_cpp.

• Tambah fail pengepala main.h:

python jax_to_ir.py \
  --fn jax_example.prog.fn \
  --input_shapes '[("x", "f32[2,2]"), ("y", "f32[2,2")]' \
  --constants '{"z": 2.0}' \
  --ir_format HLO \
  --ir_human_dest /tmp/fn_hlo.txt \
  --ir_dest /tmp/fn_hlo.pb

• Kemas kini fail BUILD untuk mencipta pustaka kongsi:

#ifndef MAIN_H  
#define MAIN_H  

extern "C" {  
  int bar(int foo);  
}  

#endif  

Kompilasi dengan Bazel:

cc_shared_library(  
   name = "jax",  
   deps = [":main"],  
   visibility = ["//visibility:public"],  
)  

Anda akan menemui libjax.dylib yang dikompilasi dalam direktori output.

3. Sambungkan Dart dengan C Menggunakan FFI

Gunakan pakej FFI Dart untuk berkomunikasi dengan perpustakaan C. Cipta fail jax.dart:

bazel build examples/jax_cpp:jax  

Sertakan perpustakaan dinamik dalam direktori projek anda. Uji dengan:

import 'dart:ffi';  
import 'package:dynamic_library/dynamic_library.dart';  

typedef FooCFunc = Int32 Function(Int32 bar);  
typedef FooDartFunc = int Function(int bar);  

class JAX {  
  late final DynamicLibrary dylib;  

  JAX() {  
    dylib = loadDynamicLibrary(libraryName: 'jax');  
  }  

  Function get _bar => dylib.lookupFunction<FooCFunc, FooDartFunc>('bar');  

  int bar(int foo) {  
    return _bar(foo);  
  }  
}  

Anda akan melihat output daripada pustaka C dalam konsol anda.

? Langkah Seterusnya

Dengan persediaan ini, anda boleh:

• Optimumkan model ML dengan JAX dan XLA.
• Jalankan algoritma yang berkuasa secara tempatan.

Kes penggunaan yang berpotensi termasuk:

• Algoritma carian (cth., A*).
• Pengoptimuman gabungan (cth., penjadualan).
• Pemprosesan imej (cth., pengesanan tepi).

JAX merapatkan jurang antara pembangunan berasaskan Python dan prestasi peringkat pengeluaran, membenarkan jurutera ML menumpukan pada algoritma tanpa perlu risau tentang kod C peringkat rendah.

Kami sedang membina platform AI yang canggih dengan token sembang tanpa had dan ingatan jangka panjang, memastikan interaksi yang lancar dan sedar konteks yang berkembang dari semasa ke semasa.

Ia percuma sepenuhnya dan anda boleh mencubanya dalam IDE semasa anda juga.

Atas ialah kandungan terperinci Menjalankan Program JAX daripada Dart Menggunakan C FFI. Untuk maklumat lanjut, sila ikut artikel berkaitan lain di laman web China PHP!