Tensorflow学习笔记（三）下 -- @tf.function的使用与自定义求导
继续接着上一篇来看：
import matplotlib as mpl
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
import sklearn
import pandas as pd
import os
import sys
import time
import tensorflow as tf
from tensorflow import keras
print(tf.__version__)
print(sys.version_info)
for module in mpl, np, pd, sklearn, tf, keras:
print(module.__name__, module.__version__)
# tf.function and auto-graph.
def scaled_elu(z, scale=1.0, alpha=1.0):
# z >= 0 ? scale * z : scale * alpha * tf.nn.elu(z)
is_positive = tf.greater_equal(z, 0.0)
return scale * tf.where(is_positive, z, alpha * tf.nn.elu(z))
print(scaled_elu(tf.constant(-3.)))
print(scaled_elu(tf.constant([-3., -2.5])))
scaled_elu_tf = tf.function(scaled_elu)
print(scaled_elu_tf(tf.constant(-3.)))
print(scaled_elu_tf(tf.constant([-3., -2.5])))
print(scaled_elu_tf.python_function is scaled_elu)
'''
tf.Tensor(-0.95021296, shape=(), dtype=float32)
tf.Tensor([-0.95021296 -0.917915 ], shape=(2,), dtype=float32)
tf.Tensor(-0.95021296, shape=(), dtype=float32)
tf.Tensor([-0.95021296 -0.917915 ], shape=(2,), dtype=float32)
True
'''
%timeit scaled_elu(tf.random.normal((1000, 1000)))
%timeit scaled_elu_tf(tf.random.normal((1000, 1000)))
'''
25.4 ms ± 1.75 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
23.5 ms ± 454 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
'''
# 1 + 1/2 + 1/2^2 + ... + 1/2^n
@tf.function
def converge_to_2(n_iters):
total = tf.constant(0.)
increment = tf.constant(1.)
for _ in range(n_iters):
total += increment
increment /= 2.0
return total
print(converge_to_2(20))
# tf.Tensor(1.9999981, shape=(), dtype=float32)
def display_tf_code(func):
code = tf.autograph.to_code(func)
from IPython.display import display, Markdown
display(Markdown('```python\n{}\n```'.format(code)))
display_tf_code(scaled_elu)
#
var = tf.Variable(0.)
@tf.function
def add_21():
return var.assign_add(21) # +=
print(add_21())
# tf.Tensor(21.0, shape=(), dtype=float32)
@tf.function(input_signature=[tf.TensorSpec([None], tf.int32, name='x')])
def cube(z):
return tf.pow(z, 3)
try:
print(cube(tf.constant([1., 2., 3.])))
except ValueError as ex:
print(ex)
print(cube(tf.constant([1, 2, 3])))
'''
Python inputs incompatible with input_signature: inputs ((,)), input_signature ((TensorSpec(shape=(None,), dtype=tf.int32, name='x'),))
tf.Tensor([ 1 8 27], shape=(3,), dtype=int32)
'''
# @tf.function py func -> tf graph
# get_concrete_function -> add input signature -> SavedModel
cube_func_int32 = cube.get_concrete_function(
tf.TensorSpec([None], tf.int32))
print(cube_func_int32)
'''
'''
print(cube_func_int32 is cube.get_concrete_function(
tf.TensorSpec([5], tf.int32)))
print(cube_func_int32 is cube.get_concrete_function(
tf.constant([1, 2, 3])))
'''
True
True
'''
cube_func_int32.graph
#
cube_func_int32.graph.get_operations()
'''
[,
,
,
]
'''
pow_op = cube_func_int32.graph.get_operations()[2]
print(pow_op)
'''
name: "Pow"
op: "Pow"
input: "x"
input: "Pow/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
'''
print(list(pow_op.inputs))
print(list(pow_op.outputs))
'''
[, ]
[]
'''
cube_func_int32.graph.get_operation_by_name("x") #用于模型的保存
#
cube_func_int32.graph.get_tensor_by_name("x:0")
#
cube_func_int32.graph.as_graph_def()
'''
node {
name: "x"
op: "Placeholder"
attr {
key: "_user_specified_name"
value {
s: "x"
}
}
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "shape"
value {
shape {
dim {
size: -1
}
}
}
}
}
node {
name: "Pow/y"
op: "Const"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 3
}
}
}
}
node {
name: "Pow"
op: "Pow"
input: "x"
input: "Pow/y"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
node {
name: "Identity"
op: "Identity"
input: "Pow"
attr {
key: "T"
value {
type: DT_INT32
}
}
}
versions {
producer: 27
}
'''
作者：weixin_43107805
