Saat ini, terdapat tiga framework yang mendominasi dunia riset dan produksi deep learning. Masing-masing memiliki keunggulan tersendiri: satu dikenal karena kemudahan penggunaannya, satu lagi karena kelengkapan fitur dan kematangan ekosistemnya, dan satu lainnya karena kemampuan skalabilitasnya yang luar biasa. Lalu, framework mana yang sebaiknya Anda pilih?

Deep learning telah membawa perubahan besar maupun kecil dalam kehidupan kita sehari-hari. Mulai dari Siri dan Alexa yang memahami perintah suara, aplikasi penerjemah real-time di ponsel, hingga teknologi penglihatan komputer yang menggerakkan traktor pintar, robot di gudang, serta mobil otonom — setiap bulan, selalu ada inovasi baru yang bermunculan. Hampir seluruh aplikasi deep learning ini dikembangkan menggunakan salah satu dari tiga framework utama: TensorFlow, PyTorch, atau JAX.

Framework mana yang paling tepat untuk Anda gunakan? Dalam artikel ini, akan dibahas perbandingan ketiganya secara umum: TensorFlow, PyTorch, dan JAX. Pembahasan akan mencakup kekuatan masing-masing framework dalam berbagai jenis aplikasi, serta mempertimbangkan aspek lain seperti dukungan komunitas dan kemudahan penggunaan.

Mengapa Memilih TensorFlow?

Pada era 1970-an dan 1980-an, dunia komputasi punya slogan populer: “Tidak ada yang dipecat karena membeli IBM.” Di dekade 2010-an, hal serupa bisa dikatakan tentang penggunaan TensorFlow untuk deep learning. Namun seperti yang kita tahu, IBM akhirnya tersisih saat memasuki tahun 1990-an. Lalu, setelah tujuh tahun sejak dirilis pada 2015, apakah TensorFlow masih kompetitif di dekade baru ini?

Jawabannya: tentu saja. TensorFlow tidak diam di tempat. TensorFlow 1.x memang dikenal dengan pendekatan membangun grafik statis yang terasa kurang natural bagi pengguna Python. Namun sejak TensorFlow 2.x, kita bisa membangun model dengan mode “eager execution” untuk evaluasi operasi secara langsung — membuat pengalaman coding lebih mirip dengan PyTorch. Di tingkat atas, TensorFlow menawarkan Keras untuk kemudahan pengembangan, dan di tingkat bawah, ada XLA (Accelerated Linear Algebra) sebagai compiler optimasi yang meningkatkan kecepatan, terutama pada GPU. XLA juga merupakan kunci untuk memaksimalkan performa TPU (Tensor Processing Units) milik Google yang luar biasa dalam melatih model berskala besar.

Selain itu, TensorFlow juga telah lama unggul di banyak aspek. Mau menyajikan model dalam platform yang matang dan stabil? Ada TensorFlow Serving. Mau deploy model ke web, smartphone, atau perangkat IoT dengan sumber daya terbatas? TensorFlow.js dan TensorFlow Lite sudah sangat matang untuk kebutuhan tersebut. Dan jangan lupakan: hingga saat ini, 100% deployment produksi Google masih menggunakan TensorFlow — membuktikan skalabilitas dan keandalan platform ini.

“TensorFlow tetap menjadi pilihan utama untuk aplikasi produksi besar, dengan dukungan penuh terhadap perangkat keras seperti GPU dan TPU serta ekosistem yang sangat matang.”

Namun, harus diakui, belakangan ini terasa ada semacam “penurunan energi” di sekitar proyek TensorFlow. Migrasi dari TensorFlow 1.x ke 2.x, misalnya, terbukti sangat berat. Banyak perusahaan akhirnya memilih memindahkan kode mereka ke PyTorch daripada menghabiskan waktu untuk adaptasi besar-besaran. TensorFlow juga kehilangan popularitas di kalangan peneliti, yang lebih memilih fleksibilitas PyTorch — menyebabkan penurunan jumlah publikasi ilmiah berbasis TensorFlow.

Masalah lain muncul dari kisah Keras. Setelah menjadi bagian resmi dari TensorFlow, Keras belakangan kembali dipisahkan sebagai library independen dengan jadwal rilis sendiri. Meskipun perubahan ini mungkin tidak berdampak besar pada penggunaan harian, pembalikan keputusan yang begitu besar di versi minor tentu menimbulkan keraguan.

Meski demikian, TensorFlow tetaplah sebuah framework yang andal dengan ekosistem deep learning yang sangat luas. Anda bisa membangun aplikasi dan model dari skala kecil hingga besar dengan TensorFlow, dan akan bergabung bersama banyak komunitas pengguna lainnya. Hanya saja, untuk banyak pengembang dan peneliti saat ini, TensorFlow mungkin bukan lagi pilihan pertama.

Mengapa Memilih PyTorch?

Kini bukan lagi pendatang baru yang hanya membayangi TensorFlow, PyTorch telah menjadi kekuatan utama di dunia deep learning saat ini — utamanya dalam bidang riset, namun semakin banyak pula digunakan dalam aplikasi produksi. Dengan mode eager yang kini menjadi standar pengembangan baik di TensorFlow maupun PyTorch, pendekatan PyTorch yang lebih Pythonic melalui automatic differentiation (autograd) tampaknya telah memenangkan persaingan melawan konsep static graph.

Berbeda dengan TensorFlow, PyTorch tidak mengalami perubahan besar pada kode inti sejak penghapusan Variable API pada versi 0.4. (Sebelumnya, penggunaan autograd dengan tensor mengharuskan pemanggilan Variable; kini semua objek secara langsung berupa tensor.) Namun, bukan berarti perjalanan PyTorch tanpa hambatan. Misalnya, bagi pengguna yang telah melatih model di beberapa GPU, mungkin pernah menghadapi kebingungan antara penggunaan DataParallel dan DistributedDataParallel. Meskipun sebenarnya DistributedDataParallel sangat disarankan untuk hampir semua kasus penggunaan, DataParallel tetap tersedia dan belum resmi dihentikan.

Meskipun sebelumnya PyTorch tertinggal di belakang TensorFlow dan JAX dalam hal dukungan terhadap XLA/TPU, situasinya telah membaik secara signifikan sejak tahun 2022. Kini PyTorch mendukung akses ke TPU VM maupun dukungan gaya lama untuk TPU Node, serta memungkinkan deployment ke CPU, GPU, atau TPU melalui command-line tanpa perlu mengubah satu baris pun dalam kode. Jika Anda tidak ingin berurusan dengan kode boilerplate yang terkadang harus ditulis sendiri di PyTorch, Anda dapat memanfaatkan tambahan tingkat tinggi seperti PyTorch Lightning, yang memungkinkan Anda lebih fokus pada pekerjaan utama daripada harus menulis ulang training loop. Di sisi lain, meskipun pengembangan PyTorch Mobile terus berlanjut, kematangannya masih jauh dibandingkan dengan TensorFlow Lite.

“PyTorch telah memenangkan hati komunitas penelitian berkat pendekatannya yang lebih fleksibel dan Pythonic, serta kemudahan dalam eksperimen dengan model-model baru.”

Untuk kebutuhan produksi, PyTorch kini terintegrasi dengan berbagai platform yang bersifat framework-agnostic seperti Kubeflow. Selain itu, proyek TorchServe menangani detail deployment seperti scaling, monitoring metrik, dan batch inference, menawarkan solusi MLOps yang ringkas dan langsung dikelola oleh pengembang PyTorch sendiri. Apakah PyTorch mampu menangani skala besar? Meta (sebelumnya Facebook) telah menggunakan PyTorch dalam produksi selama bertahun-tahun, sehingga klaim bahwa PyTorch tidak mampu menangani beban kerja berskala besar dapat dipastikan tidak benar. Meski demikian, dapat dikatakan bahwa untuk pelatihan model yang sangat besar — melibatkan ribuan GPU atau TPU — PyTorch mungkin masih sedikit kalah ramah dibandingkan dengan JAX.

Akhirnya, ada satu faktor besar yang tak bisa diabaikan. Popularitas PyTorch dalam beberapa tahun terakhir hampir pasti terkait dengan kesuksesan pustaka Transformers dari Hugging Face. Memang, saat ini Transformers mendukung TensorFlow dan JAX juga, tetapi awalnya dibangun berbasis PyTorch dan hingga kini tetap sangat erat kaitannya dengan framework ini. Dengan meningkatnya dominasi arsitektur Transformer, fleksibilitas PyTorch dalam riset, dan kemudahan untuk mengakses berbagai model baru dalam hitungan hari atau bahkan jam setelah publikasi melalui model hub Hugging Face, tidak heran jika PyTorch kini semakin luas digunakan di berbagai kalangan.

Mengapa Memilih JAX?

Jika Anda kurang tertarik dengan TensorFlow, Google mungkin memiliki alternatif lain untuk Anda. Setidaknya, semacam itu. JAX adalah framework deep learning yang dikembangkan, dikelola, dan digunakan oleh Google, meskipun secara resmi bukan merupakan produk Google. Namun, jika Anda menelusuri publikasi dan rilis terbaru dari Google atau DeepMind selama setahun terakhir, Anda akan menyadari bahwa banyak riset mereka kini berpindah ke JAX. Jadi, meskipun bukan “produk resmi” Google, JAX adalah alat yang digunakan para peneliti Google untuk mendorong batasan inovasi.

Apa sebenarnya JAX itu? Cara termudah untuk membayangkannya adalah: bayangkan sebuah versi NumPy yang dipercepat dengan GPU atau TPU, yang mampu secara ajaib memvektorisasi fungsi Python dan menghitung seluruh turunan fungsi tersebut secara otomatis. Selain itu, JAX dilengkapi dengan komponen JIT (Just-In-Time) yang mengoptimalkan kode Anda untuk XLA compiler, menghasilkan peningkatan performa yang signifikan dibandingkan TensorFlow dan PyTorch. Tidak jarang eksekusi kode menjadi empat hingga lima kali lebih cepat hanya dengan mengimplementasikannya ulang menggunakan JAX, tanpa perlu melakukan optimasi khusus.

“JAX mengubah cara kita melihat komputasi numerik, dengan memberikan kecepatan dan efisiensi yang luar biasa melalui akselerasi GPU/TPU dan optimisasi Just-In-Time.”

Karena JAX bekerja di tingkat yang serupa dengan NumPy, penulisan kode JAX terjadi di tingkat yang lebih rendah dibandingkan TensorFlow/Keras, bahkan dibandingkan PyTorch. Beruntung, ekosistem proyek pendukung JAX terus berkembang. Jika Anda membutuhkan pustaka neural network, ada Flax dari Google dan Haiku dari DeepMind. Untuk kebutuhan optimasi, ada Optax, sementara PIX tersedia untuk pemrosesan citra, dan masih banyak lagi. Setelah Anda terbiasa menggunakan pustaka seperti Flax, membangun jaringan saraf di JAX menjadi jauh lebih mudah. Namun, tetap perlu diingat bahwa JAX masih memiliki sejumlah kekhasan; misalnya, cara penanganan angka acak di JAX berbeda dari framework lain dan memerlukan pemahaman tambahan.

Perlukah Anda langsung beralih sepenuhnya ke JAX dan mengikuti gelombang teknologi terbaru? Mungkin iya, jika Anda mendalami riset yang melibatkan model berskala besar yang membutuhkan sumber daya luar biasa untuk dilatih. Kemajuan yang ditawarkan JAX di bidang pelatihan deterministik, serta kemampuannya dalam menangani ribuan pod TPU secara efisien, mungkin sudah cukup menjadi alasan untuk mempertimbangkan transisi.