TensorFlow 2.0 - 自定义模型、训练过程

2021-02-19 14:43:59 浏览数 (1)

文章目录

    • 1. 自定义模型
    • 2. 学习流程

学习于:简单粗暴 TensorFlow 2

1. 自定义模型

  • 重载 call() 方法,pytorch 是重载 forward() 方法
代码语言:javascript复制
import tensorflow as tf
X = tf.constant([[1.0, 2.0, 3.0],[4.0, 5.0, 6.0]])
y = tf.constant([[10.0],[20.0]])

class Linear(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.dense = tf.keras.layers.Dense(
            units=1,
            activation=None,
            kernel_initializer=tf.zeros_initializer(),
            bias_initializer=tf.zeros_initializer()
        )
    def call(self, input): # 重载 call 方法
        output = self.dense(input)
        return output

model = Linear()

# 优化器
optimizer = tf.keras.optimizers.SGD(learning_rate=0.001)

for i in range(100):
    with tf.GradientTape() as tape: # 梯度记录器
        y_pred = model(X)
        loss = tf.reduce_mean(tf.square(y_pred-y)) # 损失
    grads = tape.gradient(loss, model.variables) # 求导
    # 更新参数
    optimizer.apply_gradients(grads_and_vars=zip(grads,model.variables))

2. 学习流程

  • 加载手写数字数据集
代码语言:javascript复制
class MNistLoader():
    def __init__(self):
        data = tf.keras.datasets.mnist
        # 加载数据
        (self.train_data, self.train_label),(self.test_data, self.test_label) = data.load_data()
        # 扩展维度,灰度图1通道 [batch_size, 28, 28, chanels=1]
        self.train_data = np.expand_dims(self.train_data.astype(np.float32) / 255.0, axis=-1)
        self.test_data = np.expand_dims(self.test_data.astype(np.float32) / 255.0, axis=-1)
        self.train_label = self.train_label.astype(np.int32)
        self.test_label = self.test_label.astype(np.int32)
        # 样本个数
        self.num_train_data, self.num_test_data = self.train_data.shape[0], self.test_data.shape[0]
    def get_batch(self, batch_size):
        # 从训练集里随机取出 batch_size 个样本
        idx = np.random.randint(0, self.num_train_data, batch_size)
        return self.train_data[idx, :], self.train_label[idx]
  • 定义模型
代码语言:javascript复制
# 自定义多层感知机模型
class MLPmodel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        # 除第一维以外的维度展平
        self.flatten = tf.keras.layers.Flatten()
        self.dense1 = tf.keras.layers.Dense(units=100, activation='relu')
        self.dense2 = tf.keras.layers.Dense(units=10)
    def call(self, input):
        x = self.flatten(input)
        x = self.dense1(x)
        x = self.dense2(x)
        output = tf.nn.softmax(x)
        return output
  • 训练
代码语言:javascript复制
# 参数
num_epochs = 5
batch_size = 50
learning_rate = 1e-4

# 模型实例
mymodel = MLPmodel()
# 数据加载
data_loader = MNistLoader()
# adam 优化器
optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)

num_batches = int(data_loader.num_train_data//batch_size * num_epochs)
# 训练
for idx in range(num_batches):
	# 取出数据
    X,y = data_loader.get_batch(batch_size)
    with tf.GradientTape() as tape: # 梯度记录
        y_pred = mymodel(X) # 预测
        # 计算交叉熵损失
        loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
        loss = tf.reduce_mean(loss)
        print("batch {}, loss {}".format(idx, loss.numpy()))
    # 计算梯度
    grads = tape.gradient(loss, mymodel.variables)
    # 更新参数
    optimizer.apply_gradients(grads_and_vars=zip(grads, mymodel.variables))
  • 预测
代码语言:javascript复制
# 评估标准
sparse_categorical_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
num_batches = int(data_loader.num_test_data // batch_size)
# 预测
for idx in range(num_batches):
	# 数据区间
    start, end = idx*batch_size, (idx 1)*batch_size
    # 放入模型,预测
    y_pred = mymodel.predict(data_loader.test_data[start : end])
	# 统计更新 预测信息    
    sparse_categorical_accuracy.update_state(y_true=data_loader.test_label[start:end],
                                             y_pred=y_pred)
print("test 准确率:{}".format(sparse_categorical_accuracy.result()))
# test 准确率:0.9455000162124634

0 人点赞