问获取模型的图形

EN

import tensorflow as tf

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(32, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy'])

(train_images, train_labels), _ = tf.keras.datasets.fashion_mnist.load_data()
train_images = train_images / 255.0

model.fit(
    train_images,
    train_labels, 
    batch_size=64,
    epochs=5)

print(model.train_tf_function._concrete_stateful_fn.graph.as_graph_def()) # Output is too large
print(*[tensor for op in model.train_tf_function._concrete_stateful_fn.graph.get_operations() for tensor in op.values()], sep="\n")

Tensor("iterator:0", shape=(), dtype=resource)
Tensor("iterator_1:0", shape=(), dtype=variant)
Tensor("IteratorGetNext:0", shape=(None, 28, 28), dtype=float32)
Tensor("IteratorGetNext:1", shape=(None,), dtype=uint8)
Tensor("sequential_21/flatten_6/Const:0", shape=(2,), dtype=int32)
Tensor("sequential_21/flatten_6/Reshape:0", shape=(None, 784), dtype=float32)
Tensor("sequential_21/dense_27/MatMul/ReadVariableOp/resource:0", shape=(), dtype=resource)
Tensor("sequential_21/dense_27/MatMul/ReadVariableOp:0", shape=(784, 32), dtype=float32)
Tensor("sequential_21/dense_27/MatMul:0", shape=(None, 32), dtype=float32)
Tensor("sequential_21/dense_27/BiasAdd/ReadVariableOp/resource:0", shape=(), dtype=resource)
Tensor("sequential_21/dense_27/BiasAdd/ReadVariableOp:0", shape=(32,), dtype=float32)
Tensor("sequential_21/dense_27/BiasAdd:0", shape=(None, 32), dtype=float32)
Tensor("sequential_21/dense_27/Relu:0", shape=(None, 32), dtype=float32)
Tensor("sequential_21/dropout_31/dropout/Const:0", shape=(), dtype=float32)
Tensor("sequential_21/dropout_31/dropout/Mul:0", shape=(None, 32), dtype=float32)
Tensor("sequential_21/dropout_31/dropout/Shape:0", shape=(2,), dtype=int32)
Tensor("sequential_21/dropout_31/dropout/random_uniform/RandomUniform:0", shape=(None, 32), dtype=float32)
Tensor("sequential_21/dropout_31/dropout/GreaterEqual/y:0", shape=(), dtype=float32)
Tensor("sequential_21/dropout_31/dropout/GreaterEqual:0", shape=(None, 32), dtype=bool)
Tensor("sequential_21/dropout_31/dropout/Cast:0", shape=(None, 32), dtype=float32)
Tensor("sequential_21/dropout_31/dropout/Mul_1:0", shape=(None, 32), dtype=float32)
Tensor("sequential_21/dense_28/MatMul/ReadVariableOp/resource:0", shape=(), dtype=resource)
Tensor("sequential_21/dense_28/MatMul/ReadVariableOp:0", shape=(32, 10), dtype=float32)
Tensor("sequential_21/dense_28/MatMul:0", shape=(None, 10), dtype=float32)
Tensor("sequential_21/dense_28/BiasAdd/ReadVariableOp/resource:0", shape=(), dtype=resource)
Tensor("sequential_21/dense_28/BiasAdd/ReadVariableOp:0", shape=(10,), dtype=float32)
Tensor("sequential_21/dense_28/BiasAdd:0", shape=(None, 10), dtype=float32)
Tensor("sequential_21/dense_28/Softmax:0", shape=(None, 10), dtype=float32)
Tensor("ExpandDims/dim:0", shape=(), dtype=int32)
Tensor("ExpandDims:0", shape=(None, 1), dtype=uint8)
Tensor("sparse_categorical_crossentropy/Cast:0", shape=(None, 1), dtype=int64)
Tensor("sparse_categorical_crossentropy/Shape:0", shape=(2,), dtype=int32)
Tensor("sparse_categorical_crossentropy/Reshape/shape:0", shape=(1,), dtype=int32)
Tensor("sparse_categorical_crossentropy/Reshape:0", shape=(None,), dtype=int64)
Tensor("sparse_categorical_crossentropy/strided_slice/stack:0", shape=(1,), dtype=int32)
Tensor("sparse_categorical_crossentropy/strided_slice/stack_1:0", shape=(1,), dtype=int32)
Tensor("sparse_categorical_crossentropy/strided_slice/stack_2:0", shape=(1,), dtype=int32)
Tensor("sparse_categorical_crossentropy/strided_slice:0", shape=(), dtype=int32)
Tensor("sparse_categorical_crossentropy/Reshape_1/shape/0:0", shape=(), dtype=int32)
Tensor("sparse_categorical_crossentropy/Reshape_1/shape:0", shape=(2,), dtype=int32)
Tensor("sparse_categorical_crossentropy/Reshape_1:0", shape=(None, 10), dtype=float32)
Tensor("sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/Shape:0", shape=(1,), dtype=int32)
Tensor("sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits:0", shape=(None,), dtype=float32)
Tensor("sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits:1", shape=(None, 10), dtype=float32)
Tensor("sparse_categorical_crossentropy/weighted_loss/Const:0", shape=(), dtype=float32)
Tensor("sparse_categorical_crossentropy/weighted_loss/Mul:0", shape=(None,), dtype=float32)
Tensor("sparse_categorical_crossentropy/weighted_loss/Const_1:0", shape=(1,), dtype=int32)
Tensor("sparse_categorical_crossentropy/weighted_loss/Sum:0", shape=(), dtype=float32)
Tensor("sparse_categorical_crossentropy/weighted_loss/num_elements:0", shape=(), dtype=int32)
Tensor("sparse_categorical_crossentropy/weighted_loss/num_elements/Cast:0", shape=(), dtype=float32)
Tensor("sparse_categorical_crossentropy/weighted_loss/Rank:0", shape=(), dtype=int32)
Tensor("sparse_categorical_crossentropy/weighted_loss/range/start:0", shape=(), dtype=int32)
Tensor("sparse_categorical_crossentropy/weighted_loss/range/delta:0", shape=(), dtype=int32)
Tensor("sparse_categorical_crossentropy/weighted_loss/range:0", shape=(0,), dtype=int32)
Tensor("sparse_categorical_crossentropy/weighted_loss/Sum_1:0", shape=(), dtype=float32)
Tensor("sparse_categorical_crossentropy/weighted_loss/value:0", shape=(), dtype=float32)
Tensor("Shape:0", shape=(2,), dtype=int32)
Tensor("strided_slice/stack:0", shape=(1,), dtype=int32)
Tensor("strided_slice/stack_1:0", shape=(1,), dtype=int32)
Tensor("strided_slice/stack_2:0", shape=(1,), dtype=int32)
Tensor("strided_slice:0", shape=(), dtype=int32)
Tensor("Cast:0", shape=(), dtype=float32)
Tensor("Mul:0", shape=(), dtype=float32)
Tensor("Rank:0", shape=(), dtype=int32)
Tensor("range/start:0", shape=(), dtype=int32)
Tensor("range/delta:0", shape=(), dtype=int32)
Tensor("range:0", shape=(0,), dtype=int32)
Tensor("Sum:0", shape=(), dtype=float32)
Tensor("AssignAddVariableOp/resource:0", shape=(), dtype=resource)
Tensor("Rank_1:0", shape=(), dtype=int32)
Tensor("range_1/start:0", shape=(), dtype=int32)
Tensor("range_1/delta:0", shape=(), dtype=int32)
Tensor("range_1:0", shape=(0,), dtype=int32)
Tensor("Sum_1:0", shape=(), dtype=float32)
Tensor("AssignAddVariableOp_1/resource:0", shape=(), dtype=resource)
Tensor("Adam/gradients/ones:0", shape=(), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/value/Shape:0", shape=(0,), dtype=int32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/value/Shape_1:0", shape=(0,), dtype=int32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/value/BroadcastGradientArgs:0", shape=(None,), dtype=int32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/value/BroadcastGradientArgs:1", shape=(None,), dtype=int32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/value/div_no_nan:0", shape=(), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/value/Sum:0", shape=(), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/value/Reshape:0", shape=(), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/value/Neg:0", shape=(), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/value/div_no_nan_1:0", shape=(), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/value/div_no_nan_2:0", shape=(), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/value/mul:0", shape=(), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/value/Sum_1:0", shape=(), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/value/Reshape_1:0", shape=(), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/Reshape/shape:0", shape=(0,), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/Reshape/shape_1:0", shape=(0,), dtype=int32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/Reshape:0", shape=(), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/Const:0", shape=(0,), dtype=int32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/Tile:0", shape=(), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/Reshape_1/shape:0", shape=(1,), dtype=int32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/Reshape_1:0", shape=(1,), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/Shape:0", shape=(1,), dtype=int32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/Tile_1:0", shape=(None,), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/weighted_loss/Mul:0", shape=(None,), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/ExpandDims/dim:0", shape=(), dtype=int32)
Tensor("gradient_tape/sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/ExpandDims:0", shape=(None, 1), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/mul:0", shape=(None, 10), dtype=float32)
Tensor("gradient_tape/sparse_categorical_crossentropy/Shape:0", shape=(2,), dtype=int32)
Tensor("gradient_tape/sparse_categorical_crossentropy/Reshape:0", shape=(None, 10), dtype=float32)
Tensor("gradient_tape/sequential_21/dense_28/BiasAdd/BiasAddGrad:0", shape=(10,), dtype=float32)
Tensor("gradient_tape/sequential_21/dense_28/MatMul/MatMul:0", shape=(None, 32), dtype=float32)
Tensor("gradient_tape/sequential_21/dense_28/MatMul/MatMul_1:0", shape=(32, 10), dtype=float32)
Tensor("gradient_tape/sequential_21/dropout_31/dropout/Shape:0", shape=(2,), dtype=int32)
Tensor("gradient_tape/sequential_21/dropout_31/dropout/Shape_1:0", shape=(2,), dtype=int32)
Tensor("gradient_tape/sequential_21/dropout_31/dropout/BroadcastGradientArgs:0", shape=(None,), dtype=int32)
Tensor("gradient_tape/sequential_21/dropout_31/dropout/BroadcastGradientArgs:1", shape=(None,), dtype=int32)
Tensor("gradient_tape/sequential_21/dropout_31/dropout/Mul:0", shape=(None, 32), dtype=float32)
Tensor("gradient_tape/sequential_21/dropout_31/dropout/Sum:0", dtype=float32)
Tensor("gradient_tape/sequential_21/dropout_31/dropout/Reshape:0", shape=(None, 32), dtype=float32)
Tensor("gradient_tape/sequential_21/dropout_31/dropout/Mul_1:0", shape=(None, 32), dtype=float32)
Tensor("gradient_tape/sequential_21/dense_27/ReluGrad:0", shape=(None, 32), dtype=float32)
Tensor("gradient_tape/sequential_21/dense_27/BiasAdd/BiasAddGrad:0", shape=(32,), dtype=float32)
Tensor("gradient_tape/sequential_21/dense_27/MatMul/MatMul:0", shape=(784, 32), dtype=float32)
Tensor("Adam/Cast/ReadVariableOp/resource:0", shape=(), dtype=resource)
Tensor("Adam/Cast/ReadVariableOp:0", shape=(), dtype=float32)
Tensor("Adam/Identity:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/ReadVariableOp/resource:0", shape=(), dtype=resource)
Tensor("Adam/ReadVariableOp:0", shape=(), dtype=int64)
Tensor("Adam/add/y:0", shape=(), dtype=int64, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/add:0", shape=(), dtype=int64, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/Cast_1:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/Cast_2/ReadVariableOp/resource:0", shape=(), dtype=resource)
Tensor("Adam/Cast_2/ReadVariableOp:0", shape=(), dtype=float32)
Tensor("Adam/Identity_1:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/Cast_3/ReadVariableOp/resource:0", shape=(), dtype=resource)
Tensor("Adam/Cast_3/ReadVariableOp:0", shape=(), dtype=float32)
Tensor("Adam/Identity_2:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/Pow:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/Pow_1:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/sub/x:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/sub:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/Sqrt:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/sub_1/x:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/sub_1:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/truediv:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/mul:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/Const:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/sub_2/x:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/sub_2:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/sub_3/x:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/sub_3:0", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/Identity_3:0", shape=(784, 32), dtype=float32)
Tensor("Adam/Identity_4:0", shape=(32,), dtype=float32)
Tensor("Adam/Identity_5:0", shape=(32, 10), dtype=float32)
Tensor("Adam/Identity_6:0", shape=(10,), dtype=float32)
Tensor("Adam/IdentityN:0", shape=(784, 32), dtype=float32)
Tensor("Adam/IdentityN:1", shape=(32,), dtype=float32)
Tensor("Adam/IdentityN:2", shape=(32, 10), dtype=float32)
Tensor("Adam/IdentityN:3", shape=(10,), dtype=float32)
Tensor("Adam/IdentityN:4", shape=(784, 32), dtype=float32)
Tensor("Adam/IdentityN:5", shape=(32,), dtype=float32)
Tensor("Adam/IdentityN:6", shape=(32, 10), dtype=float32)
Tensor("Adam/IdentityN:7", shape=(10,), dtype=float32)
Tensor("Adam/Adam/update/ResourceApplyAdam/m:0", shape=(), dtype=resource, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/Adam/update/ResourceApplyAdam/v:0", shape=(), dtype=resource, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/Adam/update_1/ResourceApplyAdam/m:0", shape=(), dtype=resource, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/Adam/update_1/ResourceApplyAdam/v:0", shape=(), dtype=resource, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/Adam/update_2/ResourceApplyAdam/m:0", shape=(), dtype=resource, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/Adam/update_2/ResourceApplyAdam/v:0", shape=(), dtype=resource, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/Adam/update_3/ResourceApplyAdam/m:0", shape=(), dtype=resource, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/Adam/update_3/ResourceApplyAdam/v:0", shape=(), dtype=resource, device=/job:localhost/replica:0/task:0/device:GPU:0)
Tensor("Adam/Adam/Const:0", shape=(), dtype=int64)
Tensor("ExpandDims_1/dim:0", shape=(), dtype=int32)
Tensor("ExpandDims_1:0", shape=(None, 1), dtype=uint8)
Tensor("Cast_1:0", shape=(None, 1), dtype=float32)
Tensor("Squeeze:0", shape=(None,), dtype=float32)
Tensor("ArgMax/dimension:0", shape=(), dtype=int32)
Tensor("ArgMax:0", shape=(None,), dtype=int64)
Tensor("Cast_2:0", shape=(None,), dtype=float32)
Tensor("Equal:0", shape=(None,), dtype=bool)
Tensor("Cast_3:0", shape=(None,), dtype=float32)
Tensor("Const:0", shape=(1,), dtype=int32)
Tensor("Sum_2:0", shape=(), dtype=float32)
Tensor("AssignAddVariableOp_2/resource:0", shape=(), dtype=resource)
Tensor("Size:0", shape=(), dtype=int32)
Tensor("Cast_4:0", shape=(), dtype=float32)
Tensor("AssignAddVariableOp_3/resource:0", shape=(), dtype=resource)
Tensor("div_no_nan/ReadVariableOp:0", shape=(), dtype=float32)
Tensor("div_no_nan/ReadVariableOp_1:0", shape=(), dtype=float32)
Tensor("div_no_nan:0", shape=(), dtype=float32)
Tensor("Identity:0", shape=(), dtype=float32)
Tensor("div_no_nan_1/ReadVariableOp:0", shape=(), dtype=float32)
Tensor("div_no_nan_1/ReadVariableOp_1:0", shape=(), dtype=float32)
Tensor("div_no_nan_1:0", shape=(), dtype=float32)
Tensor("Identity_1:0", shape=(), dtype=float32)
Tensor("Const_1:0", shape=(), dtype=int64)
Tensor("AssignAddVariableOp_4/resource:0", shape=(), dtype=resource)
Tensor("Identity_2:0", shape=(), dtype=float32)
Tensor("Identity_3:0", shape=(), dtype=float32)

print(*[tensor for op in model.train_tf_function._concrete_stateful_fn.graph.get_operations() for tensor in op.values() if tensor.name.startswith('sequential')], sep="\n")

Tensor("sequential_21/flatten_6/Const:0", shape=(2,), dtype=int32)
Tensor("sequential_21/flatten_6/Reshape:0", shape=(None, 784), dtype=float32)
Tensor("sequential_21/dense_27/MatMul/ReadVariableOp/resource:0", shape=(), dtype=resource)
Tensor("sequential_21/dense_27/MatMul/ReadVariableOp:0", shape=(784, 32), dtype=float32)
Tensor("sequential_21/dense_27/MatMul:0", shape=(None, 32), dtype=float32)
Tensor("sequential_21/dense_27/BiasAdd/ReadVariableOp/resource:0", shape=(), dtype=resource)
Tensor("sequential_21/dense_27/BiasAdd/ReadVariableOp:0", shape=(32,), dtype=float32)
Tensor("sequential_21/dense_27/BiasAdd:0", shape=(None, 32), dtype=float32)
Tensor("sequential_21/dense_27/Relu:0", shape=(None, 32), dtype=float32)
Tensor("sequential_21/dropout_31/dropout/Const:0", shape=(), dtype=float32)
Tensor("sequential_21/dropout_31/dropout/Mul:0", shape=(None, 32), dtype=float32)
Tensor("sequential_21/dropout_31/dropout/Shape:0", shape=(2,), dtype=int32)
Tensor("sequential_21/dropout_31/dropout/random_uniform/RandomUniform:0", shape=(None, 32), dtype=float32)
Tensor("sequential_21/dropout_31/dropout/GreaterEqual/y:0", shape=(), dtype=float32)
Tensor("sequential_21/dropout_31/dropout/GreaterEqual:0", shape=(None, 32), dtype=bool)
Tensor("sequential_21/dropout_31/dropout/Cast:0", shape=(None, 32), dtype=float32)
Tensor("sequential_21/dropout_31/dropout/Mul_1:0", shape=(None, 32), dtype=float32)
Tensor("sequential_21/dense_28/MatMul/ReadVariableOp/resource:0", shape=(), dtype=resource)
Tensor("sequential_21/dense_28/MatMul/ReadVariableOp:0", shape=(32, 10), dtype=float32)
Tensor("sequential_21/dense_28/MatMul:0", shape=(None, 10), dtype=float32)
Tensor("sequential_21/dense_28/BiasAdd/ReadVariableOp/resource:0", shape=(), dtype=resource)
Tensor("sequential_21/dense_28/BiasAdd/ReadVariableOp:0", shape=(10,), dtype=float32)
Tensor("sequential_21/dense_28/BiasAdd:0", shape=(None, 10), dtype=float32)
Tensor("sequential_21/dense_28/Softmax:0", shape=(None, 10), dtype=float32)