TensorFlow2.0 前向传播 张量基础的使用

TensorFlow2.0 前向传播 张量基础的使用

1.导入与设置

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import datasets
print(tf.__version__)

# 使用GPU的 设置
# 获取物理GPU的个数
gpus = tf.config.experimental.list_physical_devices("GPU")  
for gpu in gpus:
    # 设置内存自增长方式
    tf.config.experimental.set_memory_growth(gpu, True)  
print("物理GPU个数:", len(gpus))

# 获取逻辑GPU的个数
logical_gpus = tf.config.experimental.list_logical_devices("GPU") 
print("逻辑GPU个数:", len(logical_gpus))

2.加载数据与预处理

# 加载数据集
(x, y), _ = datasets.mnist.load_data()
print(x.shape)
print(y.shape)

# 类型转换
x = tf.convert_to_tensor(x, dtype = tf.float32)
y = tf.convert_to_tensor(y, dtype = tf.int32)
print(x.dtype)
print(y.dtype)

# 查看X的最小值 最大值
print(tf.reduce_min(x), tf.reduce_max(x))
print(tf.reduce_min(y), tf.reduce_max(y))

# 图像预处理
# x[0, 255] -> [0, 1]
x = x / 255.0
print(tf.reduce_min(x), tf.reduce_max(x))

3.数据集制作 变量初始化

# 制作 Dataset
train_db = tf.data.Dataset.from_tensor_slices((x, y)).batch(128)

# 查看数据集 shape
train_iter = iter(train_db)
sample = next(train_iter)
print("batch:", sample[0].shape, sample[1].shape)

# 初始化 参数 学习率
lr = 1e-3

# 随机初始化参数
# [b, 784] ==> [b, 256] ==> [b, 128] ==> [b, 10]
# [dim_in, dim_out], [dim_out]
w1 = tf.Variable(tf.random.truncated_normal([784, 256], stddev=0.1))
b1 = tf.Variable(tf.zeros([256]))
w2 = tf.Variable(tf.random.truncated_normal([256, 128], stddev=0.1))
b2 = tf.Variable(tf.zeros([128]))
w3 = tf.Variable(tf.random.truncated_normal([128, 10], stddev=0.1))
b3 = tf.Variable(tf.zeros([10]))

4.训练数据集 计算loss 计算梯度

# 循环训练(外层循环)
for epoch in range(10):
    # 数据集循环(内层循环)
    for step, (x, y) in enumerate(train_db):
        # x:[128, 28, 28]
        # y:[128]
        
        # 维度变换
        # x:[128, 28, 28] ==> [128, 28*28]
        x = tf.reshape(x, [-1, 28*28])
        
        # 记住并更新 参数变量
        with tf.GradientTape() as tape: # tf.Variable
            
            # 第一层网络 矩阵相乘 自动broadcast
            h1 = x @ w1 + tf.broadcast_to(b1, [x.shape[0], 256])
            # x:[b, 28 * 28]
            # h1 = x @ w1 + b1
            # [b, 784] @ [784, 256] + [256] ==> [b, 256] + [256] 
            # ==> [b, 256] + [b, 256]
            
            # 添加激活函数
            h1 = tf.nn.relu(h1)
            
            # 第二层网络 矩阵相乘 + 激活函数
            # [b, 256] ==> [b, 128]
            h2 = h1 @ w2 + b2
            h2 = tf.nn.relu(h2)
            
            # 第三层网络 输出层 矩阵相乘
            # [b, 128] ==> [b, 10]
            out = h2 @ w3 + b3

            # 计算误差
            # out [b, 10] 
            # y:[b]
            y_onehot = tf.one_hot(y, depth=10)

            # 均方差  mse = Mean(sum(y-out)^2)
            loss = tf.square(y_onehot - out)
            loss = tf.reduce_mean(loss)
            
        # 计算梯度 
        grads = tape.gradient(loss, [w1, b1, w2, b2, w3, b3])
        
        # 自减 更新参数变量（学习率 * 梯度变换）
        # w1 = w1 - lr * w1_grads
        w1.assign_sub(lr*grads[0])
        b1.assign_sub(lr*grads[1])

        w2.assign_sub(lr*grads[2])
        b2.assign_sub(lr*grads[3])

        w3.assign_sub(lr*grads[4])
        b3.assign_sub(lr*grads[5])
        
        # 打印部分数据
        if step % 100 == 0:
            print("epoch=", epoch, "step=", step, " loss:", float(loss))