tensorflow2按需申请显存
rtx3080兼容tf1按需申请显存
代码主要参考了
简单粗暴tensorflow2
我的tensorflow版本是:
tf-nightly-gpu 2.4.0.dev20201015
import tensorflow.compat.v1 as tf
tf.disable_eager_execution()
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.8 # 程序最多只能占用指定gpu80%的显存
config.gpu_options.allow_growth = True #程序按需申请内存
sess = tf.Session(config = config)
a = tf.get_variable(name='a', shape=[],
initializer=tf.zeros_initializer) # 指定初始化器为全0初始化
plus_one_op = tf.assign(a, a + 1.0)
sess.run(tf.global_variables_initializer()) # 初始化所有变量
for i in range(5):
sess.run(plus_one_op)
print(sess.run(a))
tf2 按需申请显存
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
tf.config.experimental.set_memory_growth(device=gpu, enable=True)
保存模型
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()
])
data_loader = MNISTLoader()
model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
loss=tf.keras.losses.sparse_categorical_crossentropy,
metrics=[tf.keras.metrics.sparse_categorical_accuracy]
)
model.fit(data_loader.train_data, data_loader.train_label, epochs=num_epochs, batch_size=batch_size)
tf.saved_model.save(model, "saved/1")
加载保存的模型
model = tf.saved_model.load("saved/1")
data_loader = MNISTLoader()
sparse_categorical_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
num_batches = int(data_loader.num_test_data // batch_size)
for batch_index in range(num_batches):
start_index, end_index = batch_index * batch_size, (batch_index + 1) * batch_size
y_pred = model(data_loader.test_data[start_index: end_index])
sparse_categorical_accuracy.update_state(y_true=data_loader.test_label[start_index: end_index], y_pred=y_pred)
print("test accuracy: %f" % sparse_categorical_accuracy.result())
保存checkpoint 断点继续训练
保留最近3次的检查点
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+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:
# 使用CheckpointManager保存模型参数到文件并自定义编号
path = manager.save(checkpoint_number=batch_index)
print("model saved to %s" % path)
继续训练
def continueTrain():
model_to_be_restored = MLP()
checkpoint = tf.train.Checkpoint(myAwesomeModel=model_to_be_restored)
checkpoint.restore(tf.train.latest_checkpoint('./save'))
# print(checkpoint)
manager = tf.train.CheckpointManager(checkpoint, directory='./save', max_to_keep=3)
optimizer = tf.keras.optimizers.Adam(learning_rate=args.learning_rate)
num_batches = int(data_loader.num_train_data // args.batch_size * args.num_epochs)
for batch_index in range(1, num_batches+1):
X, y = data_loader.get_batch(args.batch_size)
batch_index += 3300
with tf.GradientTape() as tape:
y_pred = model_to_be_restored(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_to_be_restored.variables)
optimizer.apply_gradients(grads_and_vars=zip(grads, model_to_be_restored.variables))
if batch_index % 100 == 0:
# 使用CheckpointManager保存模型参数到文件并自定义编号
path = manager.save(checkpoint_number=batch_index)
print("model saved to %s" % path)
预测
def mytest():
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))
自定义精度的写法
sparse_categorical_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
num_batches = int(test_size // args.batch_size)
# print(num_batches)
# iter 不能放在循环里面 否则只是验证了第一个batchsize的测试集
it = iter(test_dataset)
for batch_index in range(num_batches):
X,y = next(it)
y_pred = model.predict(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()))
sparse_categorical_accuracy.update_state(y_true=y, y_pred=y_pred)
print("test accuracy: %f" % sparse_categorical_accuracy.result())
损失的计算
it = iter(train_dataset.repeat())
for batch_index in range(1, num_batches + 1):
batch_index += 0
X, y = next(it)
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()))
tensorboard可视化
log_dir = 'tensorboard'
summary_writer = tf.summary.create_file_writer(log_dir) # 实例化记录器
# tensorboard 可视化
with summary_writer.as_default(): # 指定记录器
tf.summary.scalar("loss_cats_vs_dogs:", loss, step=batch_index)
# 将当前损失函数的值写入记录器