def test():
model_to_be_restored = MLP()
checkpoint = tf.train.Checkpoint(myAwesomeModel=model_to_be_restored)
checkpoint.restore(tf.train.latest_checkpoint('./save'))
y_pred = np.argmax(model_to_be_restored.predict(data_loader.test_data),
axis=-1)
print("test accuracy: %f" %
(sum(y_pred == data_loader.test_label) / data_loader.num_test_data))
def train():
model = MLP()
optimizer = tf.keras.optimizers.Adam(learning_rate=args.learning_rate)
num_batches = int(data_loader.num_train_data // args.batch_size *
args.num_epochs)
checkpoint = tf.train.Checkpoint(
myAwesomeModel=model) # 实例化Checkpoint，设置保存对象为model
for batch_index in range(1, num_batches + 1):
X, y = data_loader.get_batch(args.batch_size)
with tf.GradientTape() as tape:
y_pred = model(X)
loss = tf.keras.losses.sparse_categorical_crossentropy(
y_true=y, y_pred=y_pred)
loss = tf.reduce_mean(loss)
print("batch %d: loss %f" % (batch_index, loss.numpy()))
grads = tape.gradient(loss, model.variables)
optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
if batch_index % 100 == 0: # 每隔100个Batch保存一次
path = checkpoint.save('./save/model.ckpt') # 保存模型参数到文件
print("model saved to %s" % path)
def train():
model = MLP()
optimizer = tf.keras.optimizers.Adam(learning_rate=args.learning_rate)
num_batches = int(data_loader.num_train_data // args.batch_size *
args.num_epochs)
checkpoint = tf.train.Checkpoint(myAwesomeModel=model)
# 使用tf.train.CheckpointManager管理Checkpoint
manager = tf.train.CheckpointManager(checkpoint,
directory='./save',
max_to_keep=3)
for batch_index in range(1, num_batches):
X, y = data_loader.get_batch(args.batch_size)
with tf.GradientTape() as tape:
y_pred = model(X)
loss = tf.keras.losses.sparse_categorical_crossentropy(
y_true=y, y_pred=y_pred)
loss = tf.reduce_mean(loss)
print("batch %d: loss %f" % (batch_index, loss.numpy()))
grads = tape.gradient(loss, model.variables)
optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
if batch_index % 100 == 0:
# 使用CheckpointManager保存模型參數到文件並自定義編號
path = manager.save(checkpoint_number=batch_index)
print("model saved to %s" % path)
model = tf.keras.models.Sequential([
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(100, activation=tf.nn.relu),
tf.keras.layers.Dense(10),
tf.keras.layers.Softmax()
])
if mode == 'functional':
inputs = tf.keras.Input(shape=(28, 28, 1))
x = tf.keras.layers.Flatten()(inputs)
x = tf.keras.layers.Dense(units=100, activation=tf.nn.relu)(x)
x = tf.keras.layers.Dense(units=10)(x)
outputs = tf.keras.layers.Softmax()(x)
model = tf.keras.Model(inputs=inputs, outputs=outputs)
if mode == 'subclassing':
from zh.model.mnist.mlp import MLP
model = MLP()
if model == 'CNN':
if mode == 'sequential':
model = tf.keras.models.Sequential([
tf.keras.layers.Conv2D(
filters=32, # 卷积核数目
kernel_size=[5, 5], # 感受野大小
padding="same", # padding策略
activation=tf.nn.relu # 激活函数
),
tf.keras.layers.MaxPool2D(pool_size=[2, 2], strides=2),
tf.keras.layers.Conv2D(filters=64,
kernel_size=[5, 5],
padding="same",
activation=tf.nn.relu),
tf.keras.layers.MaxPool2D(pool_size=[2, 2], strides=2),
import tensorflow as tf
import time
from zh.model.mnist.mlp import MLP
from zh.model.utils import MNISTLoader
num_batches = 400
batch_size = 50
learning_rate = 0.001
log_dir = 'tensorboard'
data_loader = MNISTLoader()
model = MLP()
optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate)
summary_writer = tf.summary.create_file_writer(log_dir)
tf.summary.trace_on(graph=True, profiler=True)
@tf.function
def train_one_step(X, y):
with tf.GradientTape() as tape:
y_pred = model(X)
loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y,
y_pred=y_pred)
loss = tf.reduce_mean(loss)
# 注意这里使用了TensorFlow内置的tf.print()。@tf.function不支持Python内置的print方法
tf.print("loss", loss)
grads = tape.gradient(loss, model.variables)
optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))
start_time = time.time()
for batch_index in range(num_batches):