Memahami Dekorator @tf.function TensorFlow

Pengantar

Meningkatkan kinerja loop pelatihan dapat menghemat waktu komputasi berjam-jam saat melatih model pembelajaran mesin. Salah satu cara untuk meningkatkan kinerja kode TensorFlow adalah menggunakan tf.function() dekorator – perubahan satu baris sederhana yang dapat membuat fungsi Anda berjalan lebih cepat secara signifikan.

Dalam panduan singkat ini, kami akan menjelaskan caranya tf.function() meningkatkan kinerja dan melihat beberapa praktik terbaik.

Dekorator Python dan tf.fungsi()

Dalam Python, dekorator adalah fungsi yang mengubah perilaku fungsi lain. Misalnya, Anda memanggil fungsi berikut di sel buku catatan:

import tensorflow as tf

x = tf.random.uniform(shape=[100, 100], minval=-1, maxval=1, dtype=tf.dtypes.float32)

def some_costly_computation(x):
    aux = tf.eye(100, dtype=tf.dtypes.float32)
    result = tf.zeros(100, dtype = tf.dtypes.float32)
    for i in range(1,100):
        aux = tf.matmul(x,aux)/i
        result = result + aux
    return result

%timeit some_costly_computation(x)

16.2 ms ± 103 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Namun, jika kita melewatkan fungsi mahal ke a tf.function():

quicker_computation = tf.function(some_costly_computation)
%timeit quicker_computation(x)

Kita mendapatkan quicker_computation() – fungsi baru yang bekerja jauh lebih cepat dari yang sebelumnya:

4.99 ms ± 139 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Jadi, tf.function() memodifikasi some_costly_computation() dan mengeluarkan quicker_computation() fungsi. Dekorator juga memodifikasi fungsi, jadi wajar untuk membuatnya tf.function() seorang dekorator juga.

Menggunakan notasi dekorator sama dengan memanggil tf.function(function):

@tf.function
def quick_computation(x):
  aux = tf.eye(100, dtype=tf.dtypes.float32)
  result = tf.zeros(100, dtype = tf.dtypes.float32)
  for i in range(1,100):
    aux = tf.matmul(x,aux)/i
    result = result + aux
  return result

%timeit quick_computation(x)

5.09 ms ± 283 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Bagaimana tf.function() Kerja?

Kenapa kita bisa membuat fungsi tertentu berjalan 2-3x lebih cepat?

Kode TensorFlow dapat dijalankan dalam dua mode: mode bersemangat dan mode grafik. Mode Eager adalah cara standar dan interaktif untuk menjalankan kode: setiap kali Anda memanggil suatu fungsi, itu dijalankan.

Mode grafik, bagaimanapun, sedikit berbeda. Dalam mode grafik, sebelum menjalankan fungsi, TensorFlow membuat grafik komputasi, yang merupakan struktur data yang berisi operasi yang diperlukan untuk menjalankan fungsi. Grafik komputasi memungkinkan TensorFlow menyederhanakan komputasi dan menemukan peluang untuk paralelisasi. Grafik juga mengisolasi fungsi dari kode Python di atasnya, memungkinkannya dijalankan secara efisien di banyak perangkat berbeda.

Fungsi yang didekorasi dengan @tf.function dijalankan dalam dua langkah:

1. Pada langkah pertama, TensorFlow mengeksekusi kode Python untuk fungsi tersebut dan mengompilasi grafik komputasi, menunda eksekusi operasi TensorFlow apa pun.
2. Setelah itu, grafik komputasi dijalankan.

Langkah pertama akan dilewati jika tidak perlu membuat grafik komputasi baru. Ini meningkatkan kinerja fungsi tetapi juga berarti bahwa fungsi tidak akan dijalankan seperti kode Python biasa (di mana setiap baris yang dapat dieksekusi dieksekusi). Sebagai contoh, mari kita ubah fungsi kita sebelumnya:

@tf.function
def quick_computation(x):
  print('Only prints the first time!')
  aux = tf.eye(100, dtype=tf.dtypes.float32)
  result = tf.zeros(100, dtype = tf.dtypes.float32)
  for i in range(1,100):
    aux = tf.matmul(x,aux)/i
    result = result + aux
  return result

quick_computation(x)
quick_computation(x)

Ini menghasilkan:

Only prints the first time!

Grafik print() hanya dieksekusi sekali selama langkah penelusuran, yaitu saat kode Python biasa dijalankan. Panggilan berikutnya ke fungsi hanya menjalankan operasi TenforFlow dari grafik komputasi (operasi TensorFlow).

Namun, jika kita menggunakan tf.print() sebagai gantinya:

@tf.function
def quick_computation_with_print(x):
  tf.print("Prints every time!")
  aux = tf.eye(100, dtype=tf.dtypes.float32)
  result = tf.zeros(100, dtype = tf.dtypes.float32)
  for i in range(1,100):
    aux = tf.matmul(x,aux)/i
    result = result + aux
  return result

quick_computation_with_print(x)
quick_computation_with_print(x)

Prints every time!
Prints every time!

TensorFlow termasuk tf.print() dalam grafik komputasinya karena ini adalah operasi TensorFlow – bukan fungsi Python biasa.

Praktik Terbaik dengan @tf.function

Menulis Kode dengan Operasi TensorFlow

Seperti yang baru saja kita tunjukkan, beberapa bagian kode diabaikan oleh grafik komputasi. Hal ini membuat sulit untuk memprediksi perilaku fungsi saat coding dengan kode Python "normal", seperti yang baru saja kita lihat dengan print(). Sebaiknya kodekan fungsi Anda dengan operasi TensorFlow jika berlaku untuk menghindari perilaku yang tidak diharapkan.

Misalnya, for dan while loop mungkin atau mungkin tidak diubah menjadi loop TensorFlow yang setara. Oleh karena itu, lebih baik untuk menulis loop "untuk" Anda sebagai operasi vektor, jika memungkinkan. Ini akan meningkatkan kinerja kode Anda dan memastikan bahwa fungsi Anda melacak dengan benar.

Sebagai contoh, pertimbangkan hal berikut:

x = tf.random.uniform(shape=[100, 100], minval=-1, maxval=1, dtype=tf.dtypes.float32)

@tf.function
def function_with_for(x):
    summ = float(0)
    for row in x:
      summ = summ + tf.reduce_mean(row)
    return summ

@tf.function
def vectorized_function(x):
  result = tf.reduce_mean(x, axis=0)
  return tf.reduce_sum(result)


print(function_with_for(x))
print(vectorized_function(x))

%timeit function_with_for(x)
%timeit vectorized_function(x)

tf.Tensor(0.672811, shape=(), dtype=float32)
tf.Tensor(0.67281103, shape=(), dtype=float32)
1.58 ms ± 177 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
440 µs ± 8.34 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Kode dengan operasi TensorFlow jauh lebih cepat.

Hindari Referensi ke Variabel Global

Perhatikan kode berikut:

x = tf.Variable(2, dtype=tf.dtypes.float32)
y = 2

@tf.function
def power(x):
  return tf.pow(x,y)

print(power(x))

y = 3

print(power(x))

tf.Tensor(4.0, shape=(), dtype=float32)
tf.Tensor(4.0, shape=(), dtype=float32)

Pertama kali fungsi yang didekorasi power() dipanggil, nilai output adalah yang diharapkan 4. Namun, kedua kalinya, fungsi mengabaikan nilai y sudah berubah. Ini terjadi karena nilai variabel global Python dibekukan untuk fungsi setelah ditelusuri.

Cara yang lebih baik adalah dengan menggunakan tf.Variable() untuk semua variabel Anda dan berikan keduanya sebagai argumen ke fungsi Anda.

x = tf.Variable(2, dtype=tf.dtypes.float32)
y = tf.Variable(2, dtype = tf.dtypes.float32)

@tf.function
def power(x,y):
  return tf.pow(x,y)

print(power(x,y))

y.assign(3)

print(power(x,y))

tf.Tensor(4.0, shape=(), dtype=float32)
tf.Tensor(8.0, shape=(), dtype=float32)

Debugging [email dilindungi]_s

Secara umum, Anda ingin men-debug fungsi Anda dalam mode bersemangat, lalu menghiasnya dengan @tf.function setelah kode Anda berjalan dengan benar karena pesan kesalahan dalam mode bersemangat lebih informatif.

Beberapa masalah umum adalah kesalahan jenis dan kesalahan bentuk. Jenis kesalahan terjadi ketika ada ketidakcocokan dalam jenis variabel yang terlibat dalam operasi:

x = tf.Variable(1, dtype = tf.dtypes.float32)
y = tf.Variable(1, dtype = tf.dtypes.int32)

z = tf.add(x,y)

InvalidArgumentError: cannot compute AddV2 as input #1(zero-based) was expected to be a float tensor but is a int32 tensor [Op:AddV2]

Jenis kesalahan mudah merayap masuk, dan dapat dengan mudah diperbaiki dengan melemparkan variabel ke jenis yang berbeda:

y = tf.cast(y, tf.dtypes.float32)
z = tf.add(x, y) 
tf.print(z) 

Kesalahan bentuk terjadi ketika tensor Anda tidak memiliki bentuk yang dibutuhkan operasi Anda:

x = tf.random.uniform(shape=[100, 100], minval=-1, maxval=1, dtype=tf.dtypes.float32)
y = tf.random.uniform(shape=[1, 100], minval=-1, maxval=1, dtype=tf.dtypes.float32)

z = tf.matmul(x,y)

InvalidArgumentError: Matrix size-incompatible: In[0]: [100,100], In[1]: [1,100] [Op:MatMul]

Salah satu alat yang mudah digunakan untuk memperbaiki kedua jenis kesalahan adalah debugger Python interaktif, yang dapat Anda panggil secara otomatis di Notebook Jupyter menggunakan %pdb. Dengan menggunakan itu, Anda dapat mengkodekan fungsi Anda dan menjalankannya melalui beberapa kasus penggunaan umum. Jika ada kesalahan, prompt interaktif akan terbuka. Prompt ini memungkinkan Anda untuk naik dan turun lapisan abstraksi dalam kode Anda dan memeriksa nilai, jenis, dan bentuk variabel TensorFlow Anda.

Kesimpulan

Kami telah melihat bagaimana TensorFlow tf.function() membuat fungsi Anda lebih efisien, dan bagaimana @tf.function dekorator menerapkan fungsi untuk Anda sendiri.

Percepatan ini berguna dalam fungsi yang akan dipanggil berkali-kali, seperti langkah pelatihan khusus untuk model pembelajaran mesin.