基于TensorFlow2.0的Mnist手写数字分类

2024-06-15

字数统计: 412字 | 阅读时长: 2分

相比于静态图，动态图更灵活、易读，特别是在编写条件语句、循环语句时很方便，非常适合算法研究。
本文通过一个简单的Mnist手写数字识别案例，演示TensorFlow2.0中动态图的使用方法。

代码如下:

import tensorflow as tf
from tensorflow.keras.layers import Conv2D, MaxPool2D, Flatten, Dense, Dropout


class LeNet(tf.keras.Model):
    def __init__(self, classes_num):
        super(LeNet, self).__init__()

        self.conv1 = Conv2D(32, 5, activation=tf.nn.relu)
        self.pool1 = MaxPool2D(2, strides=2)
        self.conv2 = Conv2D(64, 5, activation=tf.nn.relu)
        self.pool2 = MaxPool2D(2, strides=2)

        self.flatten = Flatten()
        self.fc1 = Dense(1024, activation=tf.nn.relu)
        self.dropout = Dropout(0.5)
        self.out = Dense(classes_num)

    def call(self, x, training=False):
        x = self.pool1(self.conv1(x))
        x = self.pool2(self.conv2(x))

        x = self.flatten(x)

        x = self.fc1(x)
        x = self.dropout(x)
        x = self.out(x)

        if not training:
            x = tf.nn.softmax(x)

        return x


def cross_entropy_loss(y_pred, y_true):
    if y_true.ndim == 1:
        y_true = tf.cast(y_true, tf.int64)
        cost = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y_true, logits=y_pred)
    else:
        cost = tf.nn.softmax_cross_entropy_with_logits(labels=y_true, logits=y_pred)
    return tf.reduce_mean(cost)


def accuracy(y_pred, y_true):
    y_pred = tf.cast(y_pred, tf.int64) if y_pred.ndim == 1 else tf.argmax(y_pred, 1)
    y_true = tf.cast(y_true, tf.int64) if y_true.ndim == 1 else tf.argmax(y_true, 1)

    correct_prediction = tf.equal(y_pred, y_true)
    return tf.reduce_mean(tf.cast(correct_prediction, tf.float32), axis=-1)


def train():
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
    x_train = x_train.reshape((-1, 28, 28, 1)) / 255.
    x_test = x_test.reshape((-1, 28, 28, 1)) / 255.

    train_data = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    train_data = train_data.repeat(10).shuffle(5000).batch(32).prefetch(1)

    net = LeNet(10)
    optimizer = tf.optimizers.Adam(0.001)

    for step, (x, y_true) in enumerate(train_data, 1):
        with tf.GradientTape() as tape:
            y_pred = net(x, training=True)
            loss = cross_entropy_loss(y_pred, y_true)

        gradients = tape.gradient(loss, net.trainable_variables)
        optimizer.apply_gradients(zip(gradients, net.trainable_variables))

        if step % 4 == 0:
            acc = accuracy(y_pred, y_true)
            print(f"step:{step},loss:{loss:05f},acc:{acc:05f}")


if __name__ == '__main__':
    for dev in tf.config.experimental.list_physical_devices('GPU'):
        tf.config.experimental.set_memory_growth(dev, True)
    train()

本文作者： killf
本文链接： http://www.killf.info/机器学习/基于TensorFlow2.0的Mnist手写数字分类/
版权声明： 本博客所有文章除特别声明外，均采用 Apache License 2.0 许可协议。转载请注明出处！