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

import time
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import datasets, layers, optimizers, Sequential, metricstf.__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.biasreturn out#自定义Model，继承自Model
class MyModel(keras.Model):#需要实现__init__和call方法def __init__(self):super(MyModel, self).__init__()#自定义5层MyDense自定义Layerself.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 xcustomModel = MyModel()
customModel.build(input_shape=[None, 28*28])
customModel.summary()

运行结果：