Tensorflow模型保存与恢复

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/ZWX2445205419/article/details/89846608

打印模型文件中的变量

from tensorflow.python.tools.inspect_checkpoint import print_tensors_in_checkpoint_file
ck_dir = "checkpoints/"

print_tensors_in_checkpoint_file("checkpoints/model-9110",
                                 tensor_name="",
                                 all_tensors=True,
                                 all_tensor_names=True)

训练模型

#! -*- coding: utf-8 -*-
from tensorflow import keras
from tensorflow.contrib import layers
import tensorflow as tf
import numpy as np


mnist = keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()

x = tf.placeholder(tf.float32, [None, 28, 28, 1])
y = tf.placeholder(tf.int64, [None, ])
training = tf.placeholder(tf.bool)

max_epochs = 1000
lr = 0.001


def parse_fn(x, label):
    x = tf.reshape(x, [28, 28, 1])
    x = tf.cast(x, tf.float32)
    return x, label


train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).\
    shuffle(buffer_size=1000).\
    map(parse_fn, num_parallel_calls=4).\
    batch(batch_size=32).\
    repeat().\
    prefetch(64)
test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).\
    map(parse_fn, num_parallel_calls=4).\
    batch(batch_size=64).\
    prefetch(buffer_size=128)

train_iterator = train_dataset.make_one_shot_iterator()
test_iterator = test_dataset.make_initializable_iterator()

train_next = train_iterator.get_next()
test_next = test_iterator.get_next()

sess = tf.Session()

test_writer = tf.summary.FileWriter("logs/test")
train_writer = tf.summary.FileWriter("logs/train", graph=sess.graph)


def make_variable(name, shape, initializer, trainable=True):
    with tf.variable_scope(name):
        return tf.get_variable(name, shape, tf.float32, initializer=initializer, trainable=trainable)


def batch_normal(x, name, training, activation=tf.nn.relu):
    with tf.name_scope(name), tf.variable_scope(name):
        return layers.batch_norm(x, decay=0.9, activation_fn=activation, is_training=training, scope="batch_normal")


def max_pool(x, k, s, name, padding="SAME"):
    with tf.name_scope(name):
        return tf.nn.max_pool(x, [1, k, k, 1], [1, s, s, 1], padding, name="max_pool")


def fc(x, c_out, name, use_bias=True, activation=True):
    c_in = x.get_shape().as_list()[-1]
    with tf.name_scope(name), tf.variable_scope(name):
        weights = make_variable("weights", [c_in, c_out], initializer=tf.random_normal_initializer())
        outputs = tf.matmul(x, weights, name="matmul")
        if use_bias:
            biases = make_variable("biases", [c_out], initializer=tf.constant_initializer(0.001))
            outputs = tf.nn.bias_add(outputs, biases)
        if activation:
            outputs = tf.nn.relu(outputs, "relu")
        return outputs


def conv(x, k, s, c_in, c_out, name, use_bias=True, padding="SAME", activation=None):
    with tf.name_scope(name), tf.variable_scope(name):
        weights = make_variable("weights", [k, k, c_in, c_out], tf.random_normal_initializer())
        outputs = tf.nn.conv2d(x, weights, [1, s, s, 1], padding, name="conv")
        if use_bias:
            biases = make_variable("biases", [c_out], tf.constant_initializer(0.001))
            outputs = tf.nn.bias_add(outputs, biases)

        if activation:
            outputs = tf.nn.relu(outputs)
        return outputs


global_step = tf.Variable(0, trainable=False, name="global_step")

with tf.name_scope("conv1"):
    conv1 = conv(x, 5, 2, 1, 32, name="conv1")
    print("conv1: {}".format(conv1.shape))
    bn1 = batch_normal(conv1, "bn1", training)
    pool1 = max_pool(bn1, 2, 2, "pool1")
    print("pool1: {}".format(pool1.shape))

with tf.name_scope("conv2"):
    conv2 = conv(pool1, 3, 1, 32, 64, "conv2")
    print("conv2: {}".format(conv2.shape))
    bn2 = batch_normal(conv2, "bn2", training)
    pool2 = batch_normal(bn2, "pool_2", training)
    print("pool2: {}".format(pool2.shape))

with tf.name_scope("conv3"):
    conv3 = conv(pool2, 3, 1, 64, 128, "conv3")
    print("conv3: {}".format(conv3.shape))
    bn3 = batch_normal(conv3, "bn3", training)
    pool3 = max_pool(bn3, 2, 2, "pool3")
    print("pool3: {}".format(pool3.shape))

with tf.name_scope("flatten"):
    flatten = layers.flatten(pool3, scope="flatten")
    print("flatten: {}".format(flatten.shape))

with tf.name_scope("fc1"):
    fc1 = fc(flatten, 1024, "fc1")
    print("fc1: {}".format(fc1.shape))

with tf.name_scope("outputs"):
    outputs = fc(fc1, 10, "fc2", activation=False)
    print("outputs: {}".format(outputs.shape))

with tf.name_scope("loss"):
    cross_entroy = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y, logits=outputs)
    loss = tf.reduce_mean(cross_entroy)

with tf.name_scope("train"):
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
        train_op = tf.train.AdamOptimizer(lr).minimize(loss, global_step=global_step)

with tf.name_scope("accuracy"):
    predict_corrects = tf.equal(tf.argmax(outputs, 1), y, "corrects")
    acc = tf.reduce_mean(tf.cast(predict_corrects, tf.float32))


saver = tf.train.Saver(max_to_keep=2)

init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())


def test(sess):
    sess.run(test_iterator.initializer)
    losses = []
    accs = []
    while True:
        try:
            x_batch, y_batch = sess.run(test_next)
        except tf.errors.OutOfRangeError:
            break
        test_loss, test_acc = sess.run([loss, acc], feed_dict={x: x_batch, y: y_batch, training: False})
        losses.append(test_loss)
        accs.append(test_acc)
    ave_loss = np.mean(losses)
    ave_acc = np.mean(accs)
    summary = tf.summary.Summary(
        value=[
            tf.summary.Summary.Value(tag="test/loss", simple_value=ave_loss),
            tf.summary.Summary.Value(tag="test/acc", simple_value=ave_acc)
        ]
    )
    test_writer.add_summary(summary, global_step=global_step.eval(sess))
    print("test {}, loss: {:.3f}, acc: {:.3f}".format(global_step.eval(sess), ave_loss, ave_acc))
    return ave_loss, ave_acc


def train(sess):
    losses = []
    accs = []
    for _ in range(10):
        x_batch, y_batch = sess.run(train_next)
        _, train_loss, train_acc = sess.run([train_op, loss, acc], feed_dict={x: x_batch, y: y_batch, training: True})
        losses.append(train_loss)
        accs.append(train_acc)
    ave_loss = np.mean(losses)
    ave_acc = np.mean(accs)
    summary = tf.summary.Summary(
        value=[
            tf.summary.Summary.Value(tag="train/loss", simple_value=ave_loss),
            tf.summary.Summary.Value(tag="train/acc", simple_value=ave_acc)
        ]
    )
    test_writer.add_summary(summary, global_step=global_step.eval(sess))
    print("train {}, loss: {:.3f}, acc: {:.3f}".format(global_step.eval(sess), ave_loss, ave_acc))


best_loss = 10.0
sess.run(init_op)
for i in range(max_epochs):
    train(sess)
    if i % 10 == 0:
        loss_, acc_ = test(sess)
        if loss_ < best_loss:
            saver.save(sess, "checkpoints/model", global_step=global_step)
            best_loss = loss_

测试模型

 #! -*- coding: utf-8 -*-
from tensorflow import keras
from tensorflow.contrib import layers
import tensorflow as tf
import numpy as np


mnist = keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()

x = tf.placeholder(tf.float32, [None, 28, 28, 1])
y = tf.placeholder(tf.int64, [None, ])
training = tf.placeholder(tf.bool)

max_epochs = 1000
lr = 0.001


def parse_fn(x, label):
    x = tf.reshape(x, [28, 28, 1])
    x = tf.cast(x, tf.float32)
    return x, label


test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).\
    map(parse_fn, num_parallel_calls=4).\
    batch(batch_size=64).\
    prefetch(buffer_size=128)

test_iterator = test_dataset.make_initializable_iterator()

test_next = test_iterator.get_next()

sess = tf.Session()

global_step = tf.Variable(0, trainable=False, name="global_step")


def make_variable(name, shape, initializer, trainable=True):
    with tf.variable_scope(name):
        return tf.get_variable(name, shape, tf.float32, initializer=initializer, trainable=trainable)


def batch_normal(x, name, training, activation=tf.nn.relu):
    with tf.name_scope(name), tf.variable_scope(name):
        return layers.batch_norm(x, decay=0.9, activation_fn=activation, is_training=training, scope="batch_normal")


def max_pool(x, k, s, name, padding="SAME"):
    with tf.name_scope(name):
        return tf.nn.max_pool(x, [1, k, k, 1], [1, s, s, 1], padding, name="max_pool")


def fc(x, c_out, name, use_bias=True, activation=True):
    c_in = x.get_shape().as_list()[-1]
    with tf.name_scope(name), tf.variable_scope(name):
        weights = make_variable("weights", [c_in, c_out], initializer=tf.random_normal_initializer())
        outputs = tf.matmul(x, weights, name="matmul")
        if use_bias:
            biases = make_variable("biases", [c_out], initializer=tf.constant_initializer(0.001))
            outputs = tf.nn.bias_add(outputs, biases)
        if activation:
            outputs = tf.nn.relu(outputs, "relu")
        return outputs


def conv(x, k, s, c_in, c_out, name, use_bias=True, padding="SAME", activation=None):
    with tf.name_scope(name), tf.variable_scope(name):
        weights = make_variable("weights", [k, k, c_in, c_out], tf.random_normal_initializer())
        outputs = tf.nn.conv2d(x, weights, [1, s, s, 1], padding, name="conv")
        if use_bias:
            biases = make_variable("biases", [c_out], tf.constant_initializer(0.001))
            outputs = tf.nn.bias_add(outputs, biases)

        if activation:
            outputs = tf.nn.relu(outputs)
        return outputs


with tf.name_scope("conv1"):
    conv1 = conv(x, 5, 2, 1, 32, name="conv1")
    print("conv1: {}".format(conv1.shape))
    bn1 = batch_normal(conv1, "bn1", training)
    pool1 = max_pool(bn1, 2, 2, "pool1")
    print("pool1: {}".format(pool1.shape))

with tf.name_scope("conv2"):
    conv2 = conv(pool1, 3, 1, 32, 64, "conv2")
    print("conv2: {}".format(conv2.shape))
    bn2 = batch_normal(conv2, "bn2", training)
    pool2 = batch_normal(bn2, "pool_2", training)
    print("pool2: {}".format(pool2.shape))

with tf.name_scope("conv3"):
    conv3 = conv(pool2, 3, 1, 64, 128, "conv3")
    print("conv3: {}".format(conv3.shape))
    bn3 = batch_normal(conv3, "bn3", training)
    pool3 = max_pool(bn3, 2, 2, "pool3")
    print("pool3: {}".format(pool3.shape))

with tf.name_scope("flatten"):
    flatten = layers.flatten(pool3, scope="flatten")
    print("flatten: {}".format(flatten.shape))

with tf.name_scope("fc1"):
    fc1 = fc(flatten, 1024, "fc1")
    print("fc1: {}".format(fc1.shape))

with tf.name_scope("outputs"):
    outputs = fc(fc1, 10, "fc2", activation=False)
    print("outputs: {}".format(outputs.shape))

with tf.name_scope("loss"):
    cross_entroy = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y, logits=outputs)
    loss = tf.reduce_mean(cross_entroy)

with tf.name_scope("train"):
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
        train_op = tf.train.AdamOptimizer(lr).minimize(loss, global_step=global_step)

with tf.name_scope("accuracy"):
    predict_corrects = tf.equal(tf.argmax(outputs, 1), y, "corrects")
    acc = tf.reduce_mean(tf.cast(predict_corrects, tf.float32))


def test(sess):
    sess.run(test_iterator.initializer)
    losses = []
    accs = []
    while True:
        try:
            x_batch, y_batch = sess.run(test_next)
        except tf.errors.OutOfRangeError:
            break
        test_loss, test_acc = sess.run([loss, acc], feed_dict={x: x_batch, y: y_batch, training: False})
        losses.append(test_loss)
        accs.append(test_acc)
    ave_loss = np.mean(losses)
    ave_acc = np.mean(accs)
    print("test {}, loss: {:.3f}, acc: {:.3f}".format(global_step.eval(sess), ave_loss, ave_acc))
    return ave_loss, ave_acc


saver = tf.train.Saver()
latest_checkpoint = tf.train.latest_checkpoint("checkpoints/")
if latest_checkpoint:
    print("latest checkpint: {}".format(latest_checkpoint))
    saver.restore(sess, latest_checkpoint)

test(sess)