Tensorflow2.0笔记 - 自定义Layer和Model

本笔记主要记录如何在tensorflow中实现自定的Layer和Model。详细内容请参考代码中的链接。

import time
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import datasets, layers, optimizers, Sequential, metrics

tf.__version__
#关于自定义layer和自定义Model的相关介绍，参考下面的链接：
#https://tf.wiki/zh_hans/basic/models.html
#https://blog.csdn.net/lzs781/article/details/104741958


#自定义Dense层，继承自Layer
class MyDense(layers.Layer):
    #需要实现__init__和call方法
    def __init__(self, input_dim, output_dim):
        super(MyDense, self).__init__()
        self.kernel = self.add_weight(name='w', shape=[input_dim, output_dim], initializer=tf.random_uniform_initializer(0, 1.0))
        self.bias = self.add_weight(name='b', shape=[output_dim], initializer=tf.random_uniform_initializer(0, 1.0))

    def call(self, inputs, training=None):
        out = inputs@self.kernel + self.bias
        return out

#自定义Model，继承自Model
class MyModel(keras.Model):
    #需要实现__init__和call方法
    def __init__(self):
        super(MyModel, self).__init__()
        #自定义5层MyDense自定义Layer
        self.fc1 = MyDense(28*28, 256)
        self.fc2 = MyDense(256, 128)
        self.fc3 = MyDense(128, 64)
        self.fc4 = MyDense(64, 32)
        self.fc5 = MyDense(32, 10)

    def call(self, inputs, trainning=None):
        x = self.fc1(inputs) #会调用MyDense的call方法
        x = tf.nn.relu(x)
        x = self.fc2(x)
        x = tf.nn.relu(x)
        x = self.fc3(x)
        x = tf.nn.relu(x)
        x = self.fc4(x)
        x = tf.nn.relu(x)
        x = self.fc5(x)
        return x

customModel = MyModel()
customModel.build(input_shape=[None, 28*28])
customModel.summary()

运行结果：