tf.function的一個很酷的新功能是AutoGraph，它容許使用天然的Python語法編寫圖形代碼。python
1.tf.function裝飾器git
當使用tf.function註釋函數時，能夠像調用任何其餘函數同樣調用它。app
它將被編譯成圖，這意味着能夠得到更快執行，更好地在GPU或TPU上運行或導出到SavedModel。dom
@tf.function函數
def simple_nn_layer(x, y):oop
return tf.nn.relu(tf.matmul(x, y))性能
x = tf.random.uniform((3, 3))ui
y = tf.random.uniform((3, 3))spa
simple_nn_layer(x, y)code
array([[0.75023645, 0.19047515, 0.10737072],
[1.1521267 , 0.49491584, 0.19416495],
[0.5541876 , 0.24642248, 0.09543521]], dtype=float32)>
若是咱們檢查註釋的結果，咱們能夠看到它是一個特殊的可調用函數，它處理與TensorFlow運行時的全部交互。
simple_nn_layer
若是代碼使用多個函數，則無需對它們進行所有註釋
從帶註釋函數調用的任何函數也將以圖形模式運行。
def linear_layer(x):
return 2 * x + 1
@tf.function
def deep_net(x):
return tf.nn.relu(linear_layer(x))
deep_net(tf.constant((1, 2, 3)))
2.使用Python控制流程
在tf.function中使用依賴於數據的控制流時，可使用Python控制流語句，AutoGraph會將它們轉換爲適當的TensorFlow操做。 例如，若是語句依賴於Tensor，則語句將轉換爲tf.cond()。
@tf.function
def square_if_positive(x):
if x > 0:
x = x * x
else:
x = 0
return x
print('square_if_positive(2) = {}'.format(square_if_positive(tf.constant(2))))
print('square_if_positive(-2) = {}'.format(square_if_positive(tf.constant(-2))))
square_if_positive(2) = 4
square_if_positive(-2) = 0
AutoGraph支持常見的Python語句，例如while，if，break，continue和return，支持嵌套。 這意味着能夠在while和if語句的條件下使用Tensor表達式，或者在for循環中迭代Tensor。
@tf.function
def sum_even(items):
s = 0
for c in items:
if c % 2 > 0:
continue
s += c
return s
sum_even(tf.constant([10, 12, 15, 20]))
AutoGraph還爲高級用戶提供了低級API。 例如，咱們可使用它來查看生成的代碼。
print(tf.autograph.to_code(sum_even.python_function, experimental_optional_features=None))
from __future__ import print_function
def tf__sum_even(items):
do_return = False
retval_ = None
s = 0
def loop_body(loop_vars, s_2):
c = loop_vars
continue_ = False
cond = c % 2 > 0
def if_true():
continue_ = True
return continue_
def if_false():
return continue_
continue_ = ag__.if_stmt(cond, if_true, if_false)
cond_1 = ag__.not_(continue_)
def if_true_1():
s_1, = s_2,
s_1 += c
return s_1
def if_false_1():
return s_2
s_2 = ag__.if_stmt(cond_1, if_true_1, if_false_1)
return s_2,
s, = ag__.for_stmt(items, None, loop_body, (s,))
do_return = True
retval_ = s
return retval_
tf__sum_even.autograph_info__ = {}
一個更復雜的控制流程的例子：
@tf.function
def fizzbuzz(n):
msg = tf.constant('')
for i in tf.range(n):
if tf.equal(i % 3, 0):
msg += 'Fizz'
elif tf.equal(i % 5, 0):
msg += 'Buzz'
else:
msg += tf.as_string(i)
msg += '\n'
return msg
print(fizzbuzz(tf.constant(15)).numpy().decode())
Fizz
1
2
Fizz
4
Buzz
Fizz
7
8
Fizz
Buzz
11
Fizz
13
14
3.Keras和AutoGraph
也能夠將tf.function與對象方法一塊兒使用。 例如，能夠經過註釋模型的調用函數來裝飾自定義Keras模型。
class CustomModel(tf.keras.models.Model):
@tf.function
def call(self, input_data):
if tf.reduce_mean(input_data) > 0:
return input_data
else:
return input_data // 2
model = CustomModel()
model(tf.constant([-2, -4]))
反作用無錫人流手術多少錢 http://www.chnk120.com/
就像在eager模式下同樣，你可使用帶有反作用的操做，好比一般在tf.function中的tf.assign或tf.print，它會插入必要的控件依賴項以確保它們按順序執行。
v = tf.Variable(5)
@tf.function
def find_next_odd():
v.assign(v + 1)
if tf.equal(v % 2, 0):
v.assign(v + 1)
find_next_odd()
v
4.用AutoGraph訓練一個簡單模型
def prepare_mnist_features_and_labels(x, y):
x = tf.cast(x, tf.float32) / 255.0
y = tf.cast(y, tf.int64)
return x, y
def mnist_dataset():
(x, y), _ = tf.keras.datasets.mnist.load_data()
ds = tf.data.Dataset.from_tensor_slices((x, y))
ds = ds.map(prepare_mnist_features_and_labels)
ds = ds.take(20000).shuffle(20000).batch(100)
return ds
train_dataset = mnist_dataset()
model = tf.keras.Sequential((
tf.keras.layers.Reshape(target_shape=(28 * 28,), input_shape=(28, 28)),
tf.keras.layers.Dense(100, activation='relu'),
tf.keras.layers.Dense(100, activation='relu'),
tf.keras.layers.Dense(10)))
model.build()
optimizer = tf.keras.optimizers.Adam()
compute_loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
compute_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
def train_one_step(model, optimizer, x, y):
with tf.GradientTape() as tape:
logits = model(x)
loss = compute_loss(y, logits)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
compute_accuracy(y, logits)
return loss
@tf.function
def train(model, optimizer):
train_ds = mnist_dataset()
step = 0
loss = 0.0
accuracy = 0.0
for x, y in train_ds:
step += 1
loss = train_one_step(model, optimizer, x, y)
if tf.equal(step % 10, 0):
tf.print('Step', step, ': loss', loss, '; accuracy', compute_accuracy.result())
return step, loss, accuracy
step, loss, accuracy = train(model, optimizer)
print('Final step', step, ': loss', loss, '; accuracy', compute_accuracy.result())
Step 10 : loss 1.85892391 ; accuracy 0.37
Step 20 : loss 1.25817835 ; accuracy 0.5125
Step 30 : loss 0.871798933 ; accuracy 0.605666637
Step 40 : loss 0.669722676 ; accuracy 0.66 ...
Step 190 : loss 0.213473886 ; accuracy 0.848105252
Step 200 : loss 0.224886 ; accuracy 0.85145
Final step tf.Tensor(200, shape=(), dtype=int32) : loss tf.Tensor(0.224886, shape=(), dtype=float32) ; accuracy tf.Tensor(0.85145, shape=(), dtype=float32)
5.關於批處理的說明
在實際應用中，批處理對性能相當重要。 轉換爲AutoGraph的最佳代碼是在批處理級別決定控制流的代碼。 若是在單個示例級別作出決策，請嘗試使用批處理API來維護性能。
def square_if_positive(x):
return [i ** 2 if i > 0 else i for i in x]
square_if_positive(range(-5, 5))
[-5, -4, -3, -2, -1, 0, 1, 4, 9, 16]
# 在tensorflow中上面的代碼應該改爲下面所示
@tf.function
def square_if_positive_naive(x):
result = tf.TensorArray(tf.int32, size=x.shape[0])
for i in tf.range(x.shape[0]):
if x[i] > 0:
result = result.write(i, x[i] ** 2)
else:
result = result.write(i, x[i])
return result.stack()
square_if_positive_naive(tf.range(-5, 5))
# 也能夠怎麼寫
def square_if_positive_vectorized(x):
return tf.where(x > 0, x ** 2, x)
square_if_positive_vectorized(tf.range(-5, 5))