TF2.0 自定义层、自定义模块、自定义模型构建

目录

头文件

1、加载数据集

2、自定义层（Dense层）

3、自定义模块（设定两个Dense为一个模块）

4、自定深度神经网络模型

5、训练深度神经网络模型

---

 

头文件

import numpy as np
import tensorflow as tf
import datetime
import os

from tensorflow.keras.layers import Dense, Flatten, Conv2D, MaxPool2D

1、加载数据集

def load_dataset():
    mnist = np.load("mnist.npz")
    return mnist['x_train']/255.0, mnist['y_train'], mnist['x_test']/255.0, mnist['y_test']

x_train, y_train, x_test, y_test = load_dataset()

print (np.shape(x_train))

x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]

train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(10000).batch(32)
test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)

2、自定义层（Dense层）

class MyDense(tf.keras.layers.Layer):
    def __init__(self, units=32, **kwargs):
        super(MyDense, self).__init__(**kwargs)
        self.units = units

    def build(self, input_shape):
        super(MyDense, self).build(input_shape) # 相当于设置self.build = True
        self.w = self.add_weight(shape=(input_shape[-1], self.units),
                                 initializer='random_normal',
                                 trainable=True,
                                 name='w')
        self.b = self.add_weight(shape=(self.units, ),
                                 initializer='random_normal',
                                 trainable=True,
                                 name='b')

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

    def get_config(self):
        config = super(MyDense, self).get_config()
        config.update({'units': self.units})
        return config

3、自定义模块（设定两个Dense为一个模块）

class MyBlock(tf.keras.Model):
    def __init__(self, units=[32, 32]):
        super(MyBlock, self).__init__(name='')

        self.dense1 = MyDense(units[0])
        self.dense2 = MyDense(units[1])

    def call(self, input_tensor):
        x = self.dense1(input_tensor)
        return self.dense2(x)

4、自定深度神经网络模型

包括卷积层，最大池化层，MyBlock，以及输出层

class MyModel(tf.keras.Model):
    def __init__(self, num_classes=10):
        super(MyModel, self).__init__()

        self.conv1 = Conv2D(4, 3, activation='relu')
        self.pool1 = MaxPool2D(pool_size=(2, 2))
        self.conv2 = Conv2D(6, 3, activation='relu')
        self.pool2 = MaxPool2D(pool_size=(2, 2))
        self.flatten = Flatten()
        self.block = MyBlock(units=[60, 40])
        self.d = Dense(10)

    @tf.function(input_signature=[tf.TensorSpec([None, 28, 28, 1], tf.float32, name='inputs')])
    def call(self, input_tensor):
        x = self.conv1(input_tensor)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.pool2(x)
        x = self.flatten(x)
        x = self.block(x)
        return self.d(x)

5、训练深度神经网络模型

def training_model():
    Epochs = 10
    model = MyModel()
    loss_obj = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    optimizer = tf.keras.optimizers.Adam(0.001)

    train_loss = tf.keras.metrics.Mean(name='train_loss')
    test_loss = tf.keras.metrics.Mean(name='test_loss')
    test_acc = tf.keras.metrics.SparseCategoricalAccuracy(name='test_acc')

    def train_step(images, labels):
        with tf.GradientTape() as tape:
            logits = model(images)
            loss = loss_obj(labels, logits)

        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

        train_loss(loss)

    def test_step(images, labels):
        logits = model(images)
        loss = loss_obj(labels, logits)

        test_loss(loss)
        test_acc(labels, logits)

    for epoch in range(Epochs):
        train_loss.reset_states()
        test_loss.reset_states()
        test_acc.reset_states()

        for images, labels in train_ds:
            train_step(images, labels)

        for images, labels in test_ds:
            test_step(images, labels)


        tmp = 'Epoch {}, train_loss: {}, test_loss: {}, test_acc: {}'
        print (tmp.format(epoch+1, train_loss.result(), test_loss.result(), test_acc.result()))


training_model()

实验结果

Epoch 1, train_loss: 0.28127846121788025, test_loss: 0.11484793573617935, test_acc: 0.963699996471405
Epoch 2, train_loss: 0.11530549079179764, test_loss: 0.08538637310266495, test_acc: 0.974399983882904
Epoch 3, train_loss: 0.09448453783988953, test_loss: 0.08153647184371948, test_acc: 0.9735000133514404
Epoch 4, train_loss: 0.08528731763362885, test_loss: 0.0694243535399437, test_acc: 0.9779000282287598
Epoch 5, train_loss: 0.07687702775001526, test_loss: 0.06510015577077866, test_acc: 0.9782000184059143
Epoch 6, train_loss: 0.07208435237407684, test_loss: 0.07370095700025558, test_acc: 0.9765999913215637
Epoch 7, train_loss: 0.06807399541139603, test_loss: 0.060320284217596054, test_acc: 0.98089998960495
Epoch 8, train_loss: 0.06528966128826141, test_loss: 0.06324119865894318, test_acc: 0.9799000024795532
Epoch 9, train_loss: 0.061740703880786896, test_loss: 0.06062839925289154, test_acc: 0.9799000024795532
Epoch 10, train_loss: 0.06004393473267555, test_loss: 0.05968930572271347, test_acc: 0.9799000024795532